Organizations: Carnegie Mellon University · Shanghai Jiao Tong University · NVIDIA · University of California, Berkeley · Princeton University · Tsinghua University · Peking University
Modern GPU workloads, especially large language model (LLM) inference, suffer from kernel launch overheads and coarse synchronization that limit inter-kernel parallelism. Recent megakernel techniques fuse multiple operators into a single persistent kernel to eliminate launch gaps and expose inter-kernel parallelism, but struggle to handle dynamic shapes and data-dependent computation in real workloads. We present Event Tensor, a unified compiler abstraction for dynamic megakernels. Event Tensor encodes dependencies between tiled tasks, and enables first-class support for both shape and data-dependent dynamism. Built atop this abstraction, our Event Tensor Compiler (ETC) applies static and dynamic scheduling transformations to generate high-performance persistent kernels. Evaluations show that ETC achieves state-of-the-art LLM serving latency while significantly reducing system warmup overhead.
Figures & tables
Figure 1: Different GPU scheduling models. Kernel-by-kernel and CUDA Graph scheduling models enforce a coarse-grained sequential execution. Megakernels break operations into smaller tasks, achieving inter-kernel parallelism.
Figure 2: Event Tensor abstraction overview. A computation graph (left) is partitioned into tiled operators ( tasks ), and the Event Tensor captures fine-grained dependencies between tasks as a first-class, symbolic-shaped object, handling the primary sources of dynamism inherent to LLM serving: 1 Shape Dynamism: Tiled tensors and Event Tensors have symbolic dimensions, such as the dynamic batch size B . 2 Data-Dependent Dynamism: The MoE layer (right) details how dependencies are resolved at runtime. Data-dependent updates and triggers (yellow arrows) use runtime-computed values such as topk and exp_indptr to dynamically manage the task execution.
Figure 3: Example Event Tensor–based program.
Figure 4: Event Tensor handles shape dynamism with symbolic-shape tensors that define a template for dependency graphs. At runtime, the template is instantiated with concrete shape values (e.g., producing a 1×2 graph for batch size 1 or a 2×2 graph for batch size 2) without recompilation or repeated graph capture.
Figure 5: Event Tensor handles data-dependent dynamism by (a) A regular workload with static dependencies. (b) A data-dependent MoE workload where runtime tensors topk and exp_indptr define an irregular task graph, enabling data-dependent event updates and task triggering.
Figure 6: GEMM + Reduce-Scatter before and after static scheduling transformation. Two separate device functions are fused into a single persistent function, with explicit notify and wait calls on the Event Tensor to coordinate dependencies.
Figure 7: Notify-and-wait mechanism for static scheduling.
Figure 8: GEMM + Reduce-Scatter after dynamic scheduling transformation. Task push and pop are inserted, and task execution is dynamically coordinated by the scheduler.
Figure 9: Push-and-pop mechanism for dynamic scheduling.
Figure 10: Comparison of runtime architectures. (a) The task graph in traditional runtime executor is materialized in memory, and only tiled operators are compiled. (b) ETC compiles scheduling logic into megakernels without runtime task graph materialization.
Figure 11: Performance results of GEMM + Reduce-Scatter on 8 B200s with dynamic scheduler.
Figure 12: Performance results of All-Gather + GEMM on 8 B200s with static scheduler.
Figure 13: Performance results of MoE layer on a single B200.
Figure 14: End-to-end performance of model serving on Qwen3-30B-A3B and Qwen3-32B (Lower is better).
Method
Warmup Time (s)
# JIT Graph Capture
SGLang (JIT)
583
51
vLLM (JIT)
123
67
ETC (AOT)
35
0
Table 1: Warmup time of Qwen3-32B model serving using different graph capturing methods.
Num tokens
1
128
1024
4096
Static
1.03
1.02
1.04
1.02
Dynamic
0.95
1.06
1.08
1.03
Table 2: Relative performance of different ETC scheduling methods on MoE layer against unfused megakernel. Higher is better.
Batch Size
1
16
32
128
Static
1.09
1.06
1.07
1.06
Dynamic
0.83
0.82
0.85
0.89
Table 3: Relative performance of different ETC scheduling methods on Qwen-3-32B with TP=4 against unfused megakernel. Higher is better.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 15: End-to-end compilation pipeline in ETC.
MLP Configurations
Name
Source Model
S
H
I
MLP-1
Qwen3-8B
8192
4096
12288
MLP-2
LLaMA-3.1-8B
8192
4096
14336
MLP-3
Gemma-2-9B
8192
3584
14336
MLP-4
Gemma-2-27B
8192
4608
36864
MLP-5
Qwen3-32B
8192
5120
25600
Appendix
Table 4: Model configurations for MLP, where S = sequence length, H = hidden dim, I = intermediate size.
Figure 16: Raw kernel relative performance results of Qwen-30B-A3B on a single B200.
Figure 17: Raw kernel relative performance results of Qwen-32B on a single B200.
Figure 18: Raw kernel relative performance results of Qwen-32B on four B200s with tensor parallelism.
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.
Deep learning inference and training performance depends critically on GPU kernel efficiency. Modern compilers such as PyTorch Inductor automatically generate GPU kernels from high-level model code, but frequently underperform expert-written implementations by wide margins. Recent LLM-assisted kernel optimizers can close this gap for standalone kernels, yet treat compiled models as black boxes, generally optimizing individual standalone kernels without respecting the compiler's structural decisions or verifying the model end-to-end. We present KernelOPT, a multi-agent system that treats compiled models as structured artifacts. It preserves vendor library calls (cuBLAS, cuDNN) and exclusively targets generated Triton sub-kernels using five profiling-guided LLM agents. A four-gate verification cascade applies static validation, multi-seed correctness checking, model-level float64-fallback verification, and performance gating (γ=1.03) to filter candidates and verify the re-stitched model end-to-end. When candidates fail verification, the system preserves the compiler baseline. The system accepts PyTorch nn Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems (100 Level 1, 100 Level 2 and 50 Level 3) on NVIDIA H200, KernelOPT achieves geometric mean speedups over torch compile of 1.40× (L1), 1.15× (L2), and 1.07× (L3) across all kernels, including fallback cases. Optimized-only geomeans (excluding cases where verification gates preserve the compiler baseline) are substantially higher: 2.54× (L1: 36/100), 1.84× (L2: 23/100), and 1.37× (L3: 11/50), reflecting where the optimizer achieves meaningful leverage.
AutoMegaKernel (AMK) compiles a HuggingFace Llama-family model into a single persistent cooperative CUDA kernel that runs the whole forward pass in one launch, with no per-model hand-written CUDA. The contribution is the system, not raw speed. A frozen schedule-IR validator statically certifies deadlock-freedom and race-freedom via static graph checks (not a mechanized proof), so an unsafe agent-proposed schedule is rejected before launch: across 7,160 adversarial schedules (6,091 unsafe) it had zero false-accepts and accepted all 360 real lowerings. The same source retargets sm_80/sm_90/sm_120 from one codebase, auto-generates correct megakernels for 10 of 10 supported models, and on a real SmolLM2-135M checkpoint reproduces HuggingFace greedy decode token-for-token (perplexity match 2.5e-7). An unattended, agent-drivable autoresearch loop self-improves the megakernel over its own baseline (1.25-1.72x). A search-found int8 (W8A16) megakernel beats CUDA-graphed cuBLAS bf16 at batch-1 decode across NVIDIA's datacenter inference fleet: L4 up to 1.33x, the current-gen L40S 1.25-1.27x, A10G up to 1.08x at scale, and the consumer RTX 5090 1.19-1.23x. The ordering is not a clean function of bandwidth (the 864 GB/s L40S beats the 600 GB/s A10G); the divide is inference-class vs training-class. AMK trails cuBLAS on the high-bandwidth training-class A100/H100, where the harness localizes the cross-SM-sync bottleneck; we report the gap plainly. This is a precision-asymmetric (W8A16 vs bf16) comparison at decode position 0; the largest real checkpoint is TinyLlama-1.1B. Code and the harness: https://github.com/RightNow-AI/AutoMegaKernel