AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

Abstract: 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

• The paper introduces a statically-checked agent harness that fuses the entire model forward pass into a single persistent kernel.
• It employs a schedule IR with explicit counter-based synchronization and formal validation to ensure deadlock- and race-free execution across NVIDIA architectures.
• Experiments demonstrate significant speedups over CUDA-graphed cuBLAS in int8 inference while reproducing token outputs losslessly for Llama-family models.

AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

Introduction and Problem Formulation

AutoMegaKernel (AMK) addresses the decoding bottleneck in LLM inference, which is fundamentally bandwidth-constrained at batch size one due to full-parameter streaming through HBM for every token. Conventional eager execution launches an individual kernel per operator, introducing significant overheads from per-op kernel launch latency and activation round-trips through device memory. CUDA Graphs mitigate CPU launch overhead but cannot eliminate HBM round-trips across operator boundaries. The megakernel paradigm solves this by fusing the entire model forward pass into a single persistent cooperative kernel—eliminating inter-op boundaries and maintaining threadblock residency throughout the pass.

AMK’s key system-level advancement is not isolated kernel speed, but the establishment of static correctness guarantees and model-agnostic kernel auto-generation. The system is designed for code agents and programmatic autotuning, supporting self-retargeting to multiple NVIDIA architectures (sm_80, sm_90, sm_120) from a shared code base.

System Architecture and Static Validation

AMK introduces a multi-layer stack where the central abstraction is a statically-typed schedule IR encoding a per-SM task-DAG with explicit counter-based synchronization. All cross-task dependencies are factored through monotonic counters, and deadlock and race-freedom are enforced via a frozen validator implementing DAG acyclicity, wait-satisfiability ($1 \leq t \leq \#\text{producers}$), and all-join rules for shared counters. This validation regime is not a full mechanized proof but is empirically sound over a trusted hand base: across 7,160 adversarially constructed schedules (6,091 unsafe), there were zero false-accepts and zero real-schedule false-negatives (accepted all 360 legitimate lowerings).

The agent harness allows for autonomous, correctness-gated search over schedule configurations and kernel parameters (threading, tiling, pipelining, etc). This guarantees that dynamic edits never violate memory safety or introduce deadlocks, which is a critical missing property in prior persistent kernel work (e.g., MPK [ch2205.xxxxx], Spector et al.'s hand-built Llama-1B kernel).

Generation, Portability, and Quantization

AMK achieves full coverage for 10/10 models in the Llama family without any hand-written CUDA, including three public HuggingFace checkpoints up to TinyLlama-1.1B. The system reliably retargets between sm_80, sm_90, and sm_120 architectures, matching HuggingFace decode outputs to within $2.5 \times 10^{-7}$ perplexity error for the SmolLM2-135M checkpoint, and reproducing token sequences and logit outputs to tight tolerances.

The megakernel synthesis path is precision-agnostic and supports automatic emission of int8 and int4 weight-only quantized kernels. The int8 path is shown to be greedy-lossless while providing a 1.12× speedup per token; the int4 path offers a 2.42× reduction in weight-traffic at a documented (∼22%) accuracy degradation. Dequantization logic is folded into the GEMV, and all variants are formally correctness-gated against CPU-side reference computation.

Empirical Performance and Analysis

AMK's strongest performance result is an inference-class cuBLAS win: on the NVIDIA L4 (sm_89), an int8 (W8A16) megakernel outperforms CUDA-graphed cuBLAS bf16 by up to 1.33× for batch-1 decode; the win persists on L40S (1.25–1.27×), A10G (up to 1.08×), and consumer RTX 5090 (1.19–1.23×). The result is robust across model scales, is not a simple function of bandwidth, and its origin is a regime effect—megakernel grid-wide synchronization costs are amortized by model size in the inference-class regime but become limiting on high-bandwidth, training-class parts (A100/H100), where AMK currently trails cuBLAS.

Comparative results against PyTorch eager and CUDA-graph eager show 3.6× speedup over per-op baseline, but the single-kernel doesn't yet surpass CUDA-graph-replayed cuBLAS in bf16 (AMK is 1.13× slower in that parity case). vLLM remains faster on high-end datacenter silicon (AMK is 1.65× slower on H100 in default cudagraph mode). The core bottleneck for parity is the GEMV implementation—raising memory-level parallelism and reducing sync is indicated as the forward path. The self-tuning harness reliably, autonomously improves kernel schedules by 1.25–1.72× over untuned baselines; further gains will require integrating more competitive Layer-1 fused kernels.

Theoretical and Practical Implications

AMK concretely demonstrates the efficacy and safety of agent-driven, statically-gated schedule search for persistent kernel synthesis. By making schedule correctness a structural invariant, the system enables fully autonomous kernel evolution and model-agnostic code generation, paving the way for domain-agnostic program synthesis for specialized accelerator kernels. The agent harness’ empirical soundness (zero validator false-accepts out of thousands of adversarial permutations) represents a meaningful step towards robust automation in ML compiler stacks.

Practically, AMK enables lightweight deployment of high-efficiency, low-latency megakernels across datacenter inference fleets with minimal manual engineering, especially in resource-constrained settings. The win over cuBLAS in int8 batch-1 inference on commodity inference parts illustrates a viable path to cost-effective, low-latency LLM deployment. The static validation infrastructure and edit-surface abstraction suggest this approach can generalize to more complex model classes and other accelerators.

Limitations and Future Directions

Despite strong inference-class results, AMK trails both cuBLAS on training-class devices and vLLM in throughput-optimized scenarios. The core performance ceiling is tied to synchronization cost and suboptimal memory bandwidth utilization in the current GEMV path. The int8 win is precision-asymmetric (lower byte traffic), and bf16 parity is not achieved. The system currently operates only on position-0, empty-KV KV-cache decode steps and does not benchmark the increased bandwidth burden of long context or multi-batch scenarios; extending the platform to cover arbitrary-context decoding and streaming workloads is non-trivial.

Quantization is realized only for naive int4 (which is lossy) and lossless int8; more aggressive or calibrated schemes, as well as integration with next-generation hand-optimized fused kernels (e.g., Marlin-class), are identified as clear avenues for closing the remaining gaps. The static correctness approach is portable but currently Llama-centric; extending coverage to MoE, SWA, and other operator-rich architectures remains open.

Conclusion

AMK introduces an agent-harnessed, statically-checked system for synthesis and self-tuning of persistent megakernels for LLM inference. It ensures schedule safety by construction and automates deployment, quantization, and self-optimization across architectures. AMK robustly wins over CUDA-graphed cuBLAS bf16 for int8 inferencing on inference-class silicon and supports fully lossless token-level equivalence against HuggingFace pipelines across multiple Llama-family models. Its contribution is the integration of static validation with autonomous kernel evolution, exposing a path to provably-safe, agent-driven kernel synthesis as AI workloads diversify and accelerators grow in complexity. The system’s limitations—chiefly kernel quality and generality—are documented explicitly, and the theoretical framework and harness architecture lay a foundation for future advances in architecture-agnostic, correctness-by-construction ML compiler backends.