EPS
← All batches·2605.11581

Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference

topic: general_aitop score: 100released: 2026-05-13first surfaced: 2026-05-13arXivPDFlinked_to_results2026-05-13

Authors: Wenxin Dong, Mingqing Hu, Guanghui Yu et al.

arXiv · PDF

Summary

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

Relevance

Not related to my persona or behavioral installation work—purely a GPU kernel optimization paper for inference speed. Only relevant if I need to optimize marker-implanted model serving, which isn't part of my current research focus.

Abstract

arXiv:2605.11581v1 Announce Type: new Abstract: When large language models (LLMs) serve real-time inference in commercial online advertising systems, end-to-end latency must be strictly bounded to the millisecond range. Yet every token generated during the decode phase triggers thousands of kernel launches, and kernel launch overhead alone can account for 14.6% of end-to-end inference time. MegaKernel eliminates launch overhead and inter-operator HBM round-trips by fusing multiple operators into a single persistent kernel. However, existing MegaKernel implementations face a fundamental tension between portability and efficiency on resource-constrained GPUs such as NVIDIA Ada: hand-tuned solutions are tightly coupled to specific architectures and lack portability, while auto-compiled approaches introduce runtime dynamic scheduling whose branch penalties are unacceptable in latency-critical settings. We observe that under a fixed deployment configuration, the optimal execution path of a MegaKernel is uniquely determined, and runtime dynamic decision-making can be entirely hoisted to compile time. Building on this insight, we propose Ada-MK: (1) a three-dimensional shared-memory constraint model combined with K-dimension splitting that reduces peak shared memory usage by 50%; (2) MLIR-based fine-grained DAG offline search that solidifies the optimal execution path, completely eliminating runtime branching; and (3) a heterogeneous hybrid inference engine that embeds MegaKernel as a plugin into TensorRT-LLM, combining high-throughput Prefill with low-latency Decode. On an NVIDIA L20, Ada-MK improves single-batch throughput by up to 23.6% over vanilla TensorRT-LLM and 50.2% over vLLM, achieving positive gains across all tested scenarios--the first industrial deployment of MegaKernel in a commercial online advertising system.