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

arXiv: https://arxiv.org/abs/2605.11581 · HTML: https://mlsanity.com/paper/2605.11581/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 31 references that are in the map.
- First posted on arXiv 12 May 2026.

## Builds on (most cited first, 17 of 17)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving](https://mlsanity.com/paper/2405.04532/) (2024)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives](https://mlsanity.com/paper/2503.20313/) (2025)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [Relax: Composable Abstractions for End-to-End Dynamic Machine Learning](https://mlsanity.com/paper/2311.02103/) (2023)
- [Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs](https://mlsanity.com/paper/2510.08726/) (2025)
- [Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers](https://mlsanity.com/paper/2601.19092/) (2026)
- [Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning](https://mlsanity.com/paper/2402.02361/) (2024)
- [Hexcute: A Compiler Framework for Automating Layout Synthesis in GPU Programs](https://mlsanity.com/paper/2504.16214/) (2025)
- [MCFuser: High-Performance and Rapid Fusion of Memory-Bound Compute-Intensive Operators](https://mlsanity.com/paper/2506.22169/) (2025)

## Built on by (most cited first, 1 of 1)

- [Measuring and Reducing WebGPU Dispatch Overhead for LLM Inference](https://mlsanity.com/paper/2608.08730/) (2026)

## Abstract (arXiv)

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.