Technical Summary: Ada-MK

Problem Statement

The deployment of Large Language Models (LLMs) in latency-sensitive commercial environments, such as online advertising systems, faces a critical bottleneck: end-to-end inference latency must be strictly bounded within the millisecond range (1–5 ms). While the MegaKernel approach—fusing multiple operators into a single persistent kernel to eliminate kernel launch overhead and HBM round-trips—offers a theoretical solution, its practical application on resource-constrained NVIDIA Ada (L20) GPUs is hindered by a fundamental tension between portability and efficiency.

Existing MegaKernel implementations, such as Stanford's, are tightly coupled to Hopper/Blackwell architectures, relying on specific hardware features like TMA (Tensor Memory Accelerator) and larger shared memory (227 KB). The Ada architecture lacks native TMA support and offers only half the shared memory (128 KB), forcing a reduction in pipeline stages from 4 to 2 and incurring significant duty-cycle losses. Conversely, auto-tuning frameworks like Mirage MPK introduce runtime dynamic decision-making (e.g., if-else branches based on memory states) to handle variability. In ultra-low-latency settings, these runtime branches cause unacceptable instruction issue penalties and pipeline bubbles, preventing the system from reaching hardware peak performance.

Methodology

The authors propose Ada-MK, a framework designed to adapt MegaKernel optimization for the Ada architecture by shifting dynamic decision-making from runtime to compile time. The system operates through two tightly coupled phases:

1. Offline MegaKernel Synthesis (Phase I):

   • MLIR-Based Fine-Grained DAG Construction: Instead of treating operators as black boxes, the system uses Torch-MLIR to decompose the Transformer Decoder and LM Head into a fine-grained Directed Acyclic Graph (DAG) at the PTX instruction level. This captures precise data dependencies (RAW/WAR) and resource constraints.
   • Adaptive Shared Memory Modeling: A three-dimensional constraint model (hardware specs, model architecture, dynamic workload) is used to calculate available shared memory pages. The system employs K-dimension fine-grained splitting, halving the K-dimension tile to reduce peak shared memory demand by 50%, enabling more pipeline stages on constrained hardware.
   • Offline Search and Path Solidification: Within the feasible search space, the system performs an exhaustive offline search to identify the optimal execution trace. Unlike traditional auto-tuners that search at the operator level, Ada-MK searches at the DAG level to capture finer-grained parallelism. The optimal path is "solidified" (serialized) into the kernel code, completely eliminating runtime branching.

2. Online Warp-Specialized Execution (Phase II):

   • Heterogeneous Hybrid Engine: Ada-MK is embedded as a plugin into the TensorRT-LLM inference engine. It utilizes a heterogeneous hybrid strategy: the Prefill phase (compute-bound, large-batch) retains TensorRT-LLM's native fused operators for high throughput, while the Decode phase (IO-bound, small-batch) switches to the optimized MegaKernel engine.
   • Warp Specialization: Inside the MegaKernel, SM resources are spatially partitioned into five warp roles (Loader, Consumer, Storer, Controller, Launcher) that collaborate via shared memory to overlap memory access with computation.

Key Contributions

1. Adaptive Shared Memory Management: The paper introduces a multi-dimensional resource model and K-dimension splitting technique that reduces peak shared memory usage by 50%. This allows for the reconstruction of efficient pipelines on Ada's limited 128 KB shared memory, overcoming the pipeline stage reduction seen in prior art.
2. Fine-Grained DAG-Based Automatic Search: By leveraging MLIR lowering to construct PTX-level dependency DAGs, Ada-MK performs offline profiling to solidify the optimal execution path. This approach eliminates the runtime dynamic decision overhead (branch penalties) inherent in previous auto-tuning methods like Mirage MPK, ensuring deterministic latency.
3. Heterogeneous Hybrid Inference Engine: The authors present the first industrial deployment of a MegaKernel as a plugin within TensorRT-LLM. This design allows the system to leverage the high parallelism of native operators for the Prefill phase while utilizing the low-latency, kernel-launch-free MegaKernel for the Decode phase, avoiding the engineering costs of a full engine replacement.
4. Quantization-Aware Optimization: The framework integrates quantization (GPTQ-W4A16) with K-dimension splitting and TensorCore-aware weight reordering. This further compresses memory pressure and enables deeper compute-memory overlap.

Experimental Results

Experiments were conducted on an NVIDIA L20 GPU using Qwen3-1.7B and Qwen2.5-1.5B models (GPTQ-W4A16 quantized) across fixed short-sequence and real-world datasets (CSL, Human-eval).

• Throughput Gains: Ada-MK improves single-batch throughput by up to 23.6% over vanilla TensorRT-LLM and 50.2% over vLLM in fixed short-sequence scenarios (input=64, output=12).
• Consistency: Positive gains were observed across all tested batch sizes (1–16) and model versions. In small-batch scenarios (BS=1), the improvement over vLLM reached 50.2% and over SGLang reached 64.5% for Qwen3-1.7B.
• Real-World Performance: On the CSL and Human-eval datasets, Ada-MK maintained the highest throughput for batch sizes 1 through 8. At larger batch sizes (BS=16), the advantage over system-level schedulers like vLLM narrowed (e.g., lagging by 3.5% on CSL), indicating that while kernel-level fusion is dominant in low-latency/small-batch settings, system-level scheduling becomes more critical in high-concurrency, long-context scenarios.
• Latency Reduction: The approach reduces end-to-end latency by 10%–50% compared to vLLM, SGLang, and vanilla TensorRT-LLM across the Qwen model series.

Significance and Claims

The paper claims to present the first industrial deployment of MegaKernel in a commercial online advertising system. Its primary significance lies in resolving the portability-efficiency trade-off that previously prevented MegaKernel adoption on resource-constrained architectures like NVIDIA Ada.

By hoisting runtime dynamic decisions to compile time through offline DAG-based search, Ada-MK achieves deterministic, low-latency inference without the branch penalties of auto-tuned solutions. The authors emphasize that this "offline search, online reuse" strategy, combined with a hybrid engine architecture, provides a viable path for deploying advanced kernel fusion techniques in production environments where strict latency bounds (1–5 ms) are mandatory. The work demonstrates that MegaKernel is not limited to high-end Hopper/Blackwell hardware but can be effectively adapted to earlier generations through rigorous, adaptive optimization.