The authors address inference latency in LLMs by optimizing "MegaKernels"—fused GPU operators that eliminate the overhead of launching thousands of separate kernels during decoding. They propose Ada-MK, which uses offline search to determine the optimal execution path at compile time (removing runtime branching penalties), reduces memory usage by 50% through better memory management, and integrates MegaKernel decoding into TensorRT-LLM. On NVIDIA L20, they achieved up to 50% throughput improvement over vLLM.
Main takeaways:
- Kernel launch overhead accounts for ~15% of LLM inference time because each token triggers thousands of kernel launches
- MegaKernels fuse operators into one persistent kernel, but existing approaches either lack portability or introduce runtime branching penalties
- Ada-MK eliminates runtime branching by determining the optimal execution path offline at compile time for a fixed deployment configuration
- Reduces peak shared memory usage by 50% through better memory constraint modeling and splitting strategies
- First industrial deployment of MegaKernel, achieving 23.6% improvement over TensorRT-LLM and 50.2% over vLLM on single-batch throughput