Ada-MK: Adaptive MegaKernel Inference

• Ada-MK is an Ada-architecture–oriented MegaKernel system that fuses multiple operators into a single persistent kernel to address latency in the LLM decode phase.
• It leverages a three-dimensional shared-memory constraint model with K-dimension splitting to optimize resource usage on NVIDIA Ada GPUs under strict SMEM limits.
• Its MLIR-based DAG search and compile-time specialization eliminate runtime branching, resulting in significant throughput improvements in low-batch, short-sequence inference.

Searching arXiv for the specified Ada-MK paper and closely related MegaKernel work. [Tool call] arxiv_search(query="Ada-MK Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference", max_results=5) Ada-MK is an Ada-architecture–oriented MegaKernel system for LLM inference that converts the decode phase into a single persistent kernel whose execution path is determined entirely at compile time (Dong et al., 12 May 2026). It was designed for commercial online advertising systems in which end-to-end latency must be bounded to the millisecond range, and it combines three elements: a three-dimensional shared-memory constraint model with K-dimension splitting, an MLIR-based fine-grained DAG offline search that removes runtime branching, and a heterogeneous hybrid inference engine that embeds the MegaKernel as a plugin into TensorRT-LLM (Dong et al., 12 May 2026). The term should not be conflated with Ada-MAC, a distinct 2013 IEEE 802.15.4 MAC protocol for health monitoring; “Ada-MK” does not appear in that earlier work (Xia et al., 2013).

1. Problem setting and design objective

Ada-MK addresses a decode-phase bottleneck that arises when LLMs serve real-time inference in online advertising systems. Under conventional execution, each generated token triggers thousands of kernel launches, and intermediate activations are repeatedly written to and read from HBM between operators. In the reported profiling of Qwen2.5-1.5B running under TensorRT-LLM on an NVIDIA Ada L20, the decode phase incurred 1,655,550 kernel launches for a representative workload; kernel launches alone consumed about 3.3 s, or 14.6% of end-to-end time (Dong et al., 12 May 2026). The paper characterizes this regime as small-batch and I/O-bound, with latency dominated by launch overhead and HBM “ping-pong.”

The motivating alternative is the MegaKernel paradigm. A MegaKernel fuses a sequence of operators, including layer norm, QKV GEMM, RoPE, attention, and MLP, into one persistent kernel; retains intermediate states in registers and shared memory; and assigns warps to specialized roles so that computation and data movement overlap within an SM (Dong et al., 12 May 2026). This removes both launch overhead and inter-operator HBM splice points.

The central difficulty is Ada-specific. NVIDIA Ada L20 lacks TMA hardware and provides 128 KB of on-chip shared memory, compared with 227 KB on H100. According to the reported analysis, straightforward porting of earlier MegaKernel designs to Ada often reduces pipeline depth from 4 to 2 stages, losing over 30% duty cycle, while auto-compiled approaches such as Mirage MPK introduce runtime state probing and if-else scheduling branches whose penalties are incompatible with latency-critical decode (Dong et al., 12 May 2026). Ada-MK is therefore framed as a portability-efficiency compromise resolved by compile-time specialization: under a fixed deployment configuration, the optimal execution path is uniquely determined and runtime dynamic decision-making can be hoisted to compile time.

2. Shared-memory modeling and K-dimension splitting

A defining component of Ada-MK is its three-dimensional shared-memory constraint model, which jointly considers hardware specification, model architecture, and dynamic workload to determine shared-memory budgeting and achievable pipeline depth (Dong et al., 12 May 2026). The available shared-memory pages are modeled as

$$N_{page} = \left\lfloor \frac{SMem_{max} - N_{stage} \cdot (Instr_{buf} + Semaphores + Scratch)}{Size_{page}} \right\rfloor .$$

With total page capacity $N_{page,total}$ and reservations for weights, quantization scales, and activations, the stage count is

$$N_{stage} = \left\lfloor \frac{N_{page,total} - N_{page,weight} - N_{page,scale} - N_{page,act}}{N_{page,per\,stage}} \right\rfloor .$$

For a fused operator DAG, the peak shared-memory demand is expressed as

$$S_{peak} = \max_{\tau \in time} \sum_{b \in Buffers} S_b(\tau),$$

subject to $S_{peak} \le SMem_{max}$ (Dong et al., 12 May 2026).

The major optimization is K-dimension splitting. If $S_w$ is the shared-memory footprint of the weight tile, and $S_{act}$ and $S_{out}$ are the activation and output footprints, then before splitting

$S^{(0)}_{peak} = S_{act} + S_{out} + S_w,$

whereas with splitting into $s$ segments,

$N_{page,total}$0

For $N_{page,total}$1, the weight component is halved, yielding an approximate 50% reduction in peak shared memory when weights dominate the footprint (Dong et al., 12 May 2026). The implementation loads only the sub-tile needed for the current K-slice and accumulates partial results across slices in registers. The paper states that synchronization is required only at K-slice boundaries, while intra-slice producer-consumer handoffs are managed through warp specialization and semaphore-based page states labeled Empty, Locked, and Ready.

Occupancy is modeled by

$N_{page,total}$2

and the search over tile sizes and K-split factors is constrained to satisfy both shared-memory limits and sufficient occupancy to hide latency (Dong et al., 12 May 2026). Cross-operator page reuse is also explicit: once activations are moved into registers, their shared-memory pages are reclaimed and reassigned to weights or epilogue accumulators. The paper identifies this activation-to-weight and activation-to-output reuse as essential on Ada for restoring 3–4 pipeline stages that would otherwise be lost to the 128 KB SMEM limit.

3. MLIR-based DAG search and compile-time path solidification

Ada-MK represents the fused block as a DAG $N_{page,total}$3 in which $N_{page,total}$4 contains PTX-level primitives such as Global→Shared async copies, Shared→Register ldmatrix, dequantization, mma.sync, epilogue arithmetic, and async writeback (Dong et al., 12 May 2026). The edge set is divided into data dependencies, derived from MLIR alias analysis as RAW/WAR/WAW relations, and resource dependencies, representing contention for shared-memory pages and warp-role slots.

The offline synthesizer searches over fusion boundaries, role mappings, tiling choices, K-split factor, prefetch distance, scheduling order, and parallelism configuration. The objective is a resource-constrained DAG makespan,

$N_{page,total}$5

with precedence constraints

$N_{page,total}$6

together with exclusivity constraints on page and warp-role timelines (Dong et al., 12 May 2026). In steady state, the per-stage duty cycle is approximated as

$N_{page,total}$7

The search procedure uses MLIR lowering to obtain PTX-granularity DAGs, constructs alias-based dependencies, enumerates candidate schedules and tilings, simulates page state transitions, and predicts duty cycle. The selected execution trace is then serialized into straight-line device code. The paper emphasizes that page selection, role handoffs, prefetch distance, and K-slice count are all fixed at compile time, so runtime if-else branching and dynamic probing disappear from the hot path (Dong et al., 12 May 2026).

This branch elimination is a central differentiator. Existing auto-compilation approaches are described as paying unacceptable branch penalties in latency-critical decode because scheduling choices are resolved dynamically. Ada-MK instead assumes a fixed deployment tuple comprising model architecture or version, quantization scheme, target batch-size range, input and output length ranges, KV-cache layout, and GPU architecture and specifications. Within that envelope, the optimal path is treated as uniquely determined and regenerated offline whenever the configuration changes (Dong et al., 12 May 2026). This suggests that Ada-MK exchanges online flexibility for deterministic execution.

4. Persistent kernel organization and inter-operator dataflow

The persistent MegaKernel in Ada-MK uses warp-specialized roles per SM: Loader, Consumer, Storer, and Controller or Launcher (Dong et al., 12 May 2026). Loader warps perform global-to-shared prefetches, address permutation, and lightweight early dequantization; Consumer warps execute ldmatrix, mma.sync, on-the-fly dequantization, and register accumulation; Storer warps handle reduction, writeback, and epilogue fusions; and Controller warps manage page allocation, semaphore signaling, progress coordination, and stage advancement.

Thread blocks are sized to respect both register pressure and $N_{page,total}$8, and warps are partitioned into role groups. Semaphore-like counters track page states, and Consumers spin only on Ready pages rather than stalling the entire block (Dong et al., 12 May 2026). Prefetch distance is tuned so that the Loader remains ahead of the Consumer, with the paper citing an $N_{page,total}$9 setting. Intermediate data remain in registers or shared memory across operator boundaries; only final outputs, or KV updates, are written to HBM. K/V cache loads for historical tokens are hoisted earlier in the pipeline window whenever permitted by RAW analysis, and address permutation is used to mitigate shared-memory bank conflicts.

The inter-operator path covers LN, QKV, attention, and MLP within one persistent loop over tokens and tiles. LN statistics and RMSNorm weights are prefetched asynchronously and decoupled from subsequent QKV computation; QKV GEMM dequantization is fused with ldmatrix to avoid temporary shared-memory materialization; and attention’s K/V fetch is moved into non-critical slots so that it overlaps with ongoing computation (Dong et al., 12 May 2026). The paper also notes that streaming reduction allows part of the epilogue to begin before both branches finish, exemplified by SwiGLU behavior in which the Up path can reduce while the Gate path executes.

Several tuning outcomes are reported. Consumer warps were reduced from 16 to 8 while the number of stages increased from 2 to 4, with the stated purpose of balancing compute and store latencies and suppressing stalls (Dong et al., 12 May 2026). With quantized weights and K-splitting, per-iteration shared memory shrank from 64 KB to 32 KB, freeing pages to deepen the pipeline. These are implementation-level manifestations of the shared-memory and scheduling models described earlier.

5. Hybrid engine integration and deployment model

Ada-MK is integrated into a heterogeneous hybrid inference engine by embedding the MegaKernel as a TensorRT-LLM plugin that replaces only the Transformer block while retaining TensorRT-LLM’s embeddings, KV-cache management, and peripheral operations (Dong et al., 12 May 2026). The stated purpose is to preserve business features, including prefix-constrained decoding and integrated generation-discrimination, while minimizing migration cost.

Execution is phase-adaptive. Prefill, characterized as long-sequence and compute-bound, remains on TensorRT-LLM’s fused operators because launch overhead is negligible there. Decode, characterized as short per-token compute and I/O-bound behavior, switches to the MegaKernel plugin to eliminate launch overhead and HBM round-trips (Dong et al., 12 May 2026). Both engines share the same KV cache and tensor layouts so that no data conversion is required. In this terminology, heterogeneity refers to combining a compiled TensorRT-LLM engine for Prefill and peripheral operations with a custom MegaKernel plugin for Decode in the same runtime.

The deployment flow is divided into two phases. In the offline phase, the Transformer block and LM Head are parsed into a fine-grained DAG; candidate schedules are pruned by resource constraints and profiled; and the optimal trace and parameters are serialized. In the online phase, the pre-serialized plan is loaded as a TensorRT-LLM plugin and executed as a single persistent kernel without dynamic control-flow decisions (Dong et al., 12 May 2026). The paper states that this compile-time approach is intended to satisfy a 1–5 ms target budget in industrial ad serving and that the search cost is one-off per configuration and amortized during production use.

A common misconception is that Ada-MK is simply a monolithic replacement for an existing inference stack. The described implementation is narrower: it replaces the Transformer block in Decode while deliberately preserving the surrounding TensorRT-LLM infrastructure (Dong et al., 12 May 2026). Another misconception is that portability is achieved by architecture-independent device code; the paper instead distinguishes generalized MLIR-based DAG construction and search from architecture-specific PTX sequences, cp.async emulation on Ada, bank-conflict avoidance permutations, and page-size choices.

6. Empirical results, limits, and scope

The reported evaluation uses a single NVIDIA L20 with 48 GB GDDR6, 864 GB/s, and CUDA 12.2. Models are Qwen3-1.7B and Qwen2.5-1.5B, both with GPTQ W4A16; batch sizes are 1, 2, 4, 8, and 16; workloads include fixed in64/out12 and real datasets CSL and Human-eval; and all frameworks are evaluated in offline batch mode to isolate kernel and operator efficiency (Dong et al., 12 May 2026).

Against vanilla TensorRT-LLM, Ada-MK improves throughput by up to 23.6% at batch size 1 for short sequences and is reported to be consistently faster across all tested scenarios, with typical gains of 4.0%–23.6% on short sequences and 4.0%–19.5% on real datasets at low-to-medium batch sizes (Dong et al., 12 May 2026). Against vLLM and SGLang, the reported gains for small-batch short sequences reach +50.2% and +71.9%, respectively, at batch size 1. The paper also records a counterexample: at larger batches and longer contexts, specifically CSL at batch size 16, vLLM can be ahead by 3.5%, which the authors attribute to system-level scheduling advantages in high-concurrency regimes.

Within the decode phase itself, Ada-MK reports an improvement of approximately 30% over a Stanford MegaKernel baseline, attributed to eliminating spurious dependencies and adopting auto-tuned pipeline configurations discovered by the DAG search (Dong et al., 12 May 2026). The gains are decomposed into four sources: removal of kernel-launch overhead, fewer HBM round-trips because operator handoffs occur through shared memory and registers, improved occupancy and duty cycle on Ada through K-splitting and page reuse, and elimination of runtime branching and state probing.

The paper explicitly limits the scope of these results. Advantages narrow, or can be eclipsed, in large-batch and long-context settings; shared-memory and register pressure on Ada impose tight ceilings on tile sizes and fused depth; and the solidified execution path is tied to a fixed configuration envelope, so substantial sequence-length variability may require multiple compiled variants and a meta-scheduler (Dong et al., 12 May 2026). Future directions listed by the authors include adapting the approach to larger models and newer Blackwell GPUs with TMA and larger SMEM, broadening fusion coverage, and adding lightweight runtime variant selection among precompiled plans without introducing divergence within the kernel.

In this sense, Ada-MK is best understood not as a general-purpose runtime scheduler, but as a compile-time-specialized decode engine for fixed deployment configurations on Ada-class GPUs (Dong et al., 12 May 2026). Its significance lies in showing that, under those conditions, MegaKernel execution can be made branch-free, deeply pipelined despite 128 KB SMEM constraints, and deployable within an industrial TensorRT-LLM stack.