🤖 machine learning

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

AutoMegaKernel is a statically-checked agent harness that automatically synthesizes a single persistent CUDA megakernel for HuggingFace Llama-family models, ensuring deadlock-free execution across diverse NVIDIA architectures while achieving up to 1.72x speedup over baselines on inference-class GPUs through self-improving agent loops.

Original authors: Jaber Jaber, Osama Jaber

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jaber Jaber, Osama Jaber

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to run a massive, complex factory (a Large Language Model) to produce a single product (one word of text).

The Old Way (Standard AI):
Currently, running this factory is like sending a worker to the warehouse to grab a specific tool, bringing it back to the assembly line, doing a tiny bit of work, sending the worker back to the warehouse for the next tool, and repeating this dozens of times for every single word.

The Problem: The worker spends most of their time walking back and forth (waiting for data from memory) rather than building. This is called being "bandwidth-bound." The factory is limited by how fast the worker can walk, not how fast they can build.
The Cost: Every time the worker stops to grab a tool, there's a tiny delay (launch latency). Doing this for every step adds up, making the process slow.

The New Way (AutoMegaKernel):
The authors of this paper built a system called AutoMegaKernel (AMK). Think of AMK as a revolutionary new factory manager that changes the rules entirely.

1. The "One-Shot" Factory

Instead of sending workers back and forth, AMK organizes the entire factory so that one single launch does the whole job.

The Analogy: Imagine instead of a worker fetching tools one by one, you have a team of workers stationed at every single machine. They are all connected by a direct, internal pipeline. When the "Start" button is pressed, the entire assembly line runs from start to finish in one continuous motion without anyone ever leaving the floor to go to the warehouse.
The Result: This removes the "walking time" and the "stop-and-start" delays, theoretically making the process much faster.

2. The "Safety Inspector" (The Statically-Checked Agent)

Building a factory where everyone works simultaneously is dangerous. If the workers aren't perfectly coordinated, they might crash into each other (race conditions) or wait forever for a tool that never arrives (deadlocks).

The Innovation: Usually, to make this work, you need a human expert to write the code manually, which is hard and error-prone. AMK uses an AI Agent to design the factory layout automatically.
The Safety Net: Before the factory ever turns on, a frozen, unchangeable Safety Inspector (a validator) checks the AI's blueprint. It mathematically proves that:
- No one will get stuck waiting forever.
- No two workers will try to use the same tool at the same time.
- The order of operations is perfect.
The Claim: The authors tested this inspector against 7,160 "tricky" or broken blueprints. The inspector caught every single unsafe one and rejected them before the factory started. It never let a dangerous plan through.

3. The "Universal Adapter" (Self-Retargeting)

Usually, a factory designed for a specific building (a specific computer chip) won't work in a different building.

The Magic: AMK is "self-retargeting." You write the blueprint once, and the system automatically adapts it to work perfectly on different types of computer chips (like the RTX 5090, A100, or H100) without needing a human to rewrite the code for each one.

4. The "Lightweight" Version (Quantization)

The system also figured out how to run the factory using "lighter" tools.

The Analogy: Instead of carrying heavy, full-size tools (high-precision math), the workers use lightweight, compact tools (int8 or int4 precision).
The Benefit: Because the tools are lighter, they can carry more of them at once, and the workers move faster.
- Int8: Works just as well as the heavy tools (lossless) but is 12% faster.
- Int4: Moves even faster (cutting the "walking time" by more than 2x), but the final product might be slightly less perfect (lossy), though still readable.

5. The Honest Results (The "Real Talk")

The authors are very transparent about where this system wins and where it loses.

Where it Wins: On "Inference-Class" chips (like the L4, L40S, A10G, and consumer RTX 5090), the AMK system is faster than the current industry standard (cuBLAS) when generating one word at a time. It beats the competition by 1.1x to 1.3x.
- Why? It successfully eliminated the "walking back and forth" delays that plague other systems.
Where it Loses: On the massive, high-speed "Training-Class" chips (like the A100 and H100), the current version is slower than the standard system.
- Why? The "safety coordination" between workers (synchronization) takes a tiny bit of time. On the fastest chips, this tiny delay becomes the bottleneck. The authors admit this openly: "We localized the bottleneck, and we are honest about it."

Summary

AutoMegaKernel is a system that uses an AI agent to design a "one-shot" factory for running AI models. It includes a rigorous safety inspector to ensure the design never crashes. It successfully runs on many different computer chips automatically.

The Big Win: It is currently the fastest way to generate text one word at a time on many popular consumer and datacenter chips, beating the standard tools by removing unnecessary delays.
The Catch: It isn't perfect yet. On the absolute fastest supercomputer chips, the coordination overhead still makes it slightly slower than the established giants, but the authors have proven the system is safe, correct, and capable of self-improvement.

Technical Summary: AutoMegaKernel (AMK)

Problem Statement

Single-stream autoregressive decoding at batch size 1 is fundamentally bandwidth-bound. The latency floor is determined by the ratio of total weight bytes to HBM bandwidth ( $t_{min} = \text{weight\_bytes} / \text{HBM\_bandwidth}$ ). Conventional execution frameworks (e.g., PyTorch, cuBLAS) fail to approach this floor because they launch a separate kernel for every operator, incurring significant CPU launch latency and forcing activations to round-trip through HBM between operations. While CUDA Graphs amortize launch overhead, they retain inter-op kernel boundaries and HBM round-trips.

Existing "megakernel" approaches (e.g., Mirage Persistent Kernel) or hand-optimized kernels address the latency floor by fusing the entire forward pass into a single persistent kernel. However, these systems lack trust and portability: they often omit static safety gates (risking deadlocks or race conditions), target specific models/architectures manually, and do not provide a structured interface for autonomous agents to safely optimize the system.

Methodology

AutoMegaKernel (AMK) is an autonomous, agent-driven system that compiles HuggingFace Llama-family models into a single persistent cooperative kernel without per-model hand-written CUDA. The system is structured into four layers and two autoresearch loops:

Layer 0 (The VM): A hand-written, frozen, per-architecture persistent megakernel (vm/scheduler.cu). It launches one cooperative threadblock per Streaming Multiprocessor (SM). Each block owns a local queue and walks a global topological order of tasks, executing wait → dispatch → signal instructions. It relies on monotonic counters for synchronization, avoiding locks.
Layer 1 (Instructions): ABI-conformant micro-kernels (GEMV, attention, RMSNorm, etc.) that are correctness-checked in isolation.
Layer 2 (Scheduler): The agent-facing edit surface. It imports a model, lowers it to a typed Schedule IR (a task-DAG synchronized by monotonic counters), and proposes structural configurations (tiling, fusion, SM assignment, pipelining depth) via a ScheduleConfig object.
Layer 3 (Roadmap): Placeholder for future dynamism (e.g., continuous batching).

The Safety Gate:
The core innovation is a statically-checked validator (validate() in schedule/ir.py). Before any GPU launch, the validator performs a set of static graph checks on the Schedule IR to certify:

Deadlock-freedom: Ensuring the producer-consumer DAG is acyclic, wait thresholds are satisfiable ( $1 \le t \le \#producers$ ), and per-SM queue orders are valid.
Race-freedom: Ensuring shared counters are joined correctly (all producers must finish before a consumer waits on a multi-producer counter) and that memory accesses respect happens-before relationships.
Correctness: Verifying output reachability and referential integrity.

If a schedule fails these checks, it is rejected immediately. This prevents unsafe agent-proposed schedules from causing hung GPUs.

Autoresearch Loop:
AMK includes an unattended, agent-drivable loop that iteratively proposes ScheduleConfig changes, validates them, runs correctness gates against a CPU reference, measures latency, and keeps or reverts changes. This loop autonomously improves the megakernel over its own baseline.

Key Contributions

Statically-Checked Schedule IR: A frozen validator that certifies deadlock- and race-freedom via static graph checks (not mechanized proofs) over a trusted base. It achieved zero false-accepts across 7,160 adversarially constructed schedules (6,091 unsafe) while accepting all 360 real model lowerings.
Automatic Generation Coverage: AMK auto-generates correct megakernels for 10 of 10 supported models (including real checkpoints like SmolLM2-135M and TinyLlama-1.1B) with zero hand-written CUDA. It successfully rejects unsupported variants (e.g., biased projections) or documents blind spots honestly.
Self-Retargeting: A single source codebase generates correct megakernels for sm_80 (A100), sm_90 (H100), and sm_120 (RTX 5090), matching eager execution bit-for-bit or within tight fp32 tolerances ( $\le 3.8 \times 10^{-5}$ ).
Quantized Megakernels: The system auto-generates int8 (greedy-lossless, 1.12× faster) and int4 (lossy, 2.42× reduced weight traffic) weight-only quantized kernels from the same IR path.
Agent-Drivable Compiler: The system treats correctness as a structural invariant, allowing an agent to search the schedule space safely without risking runtime hangs.
Honest Performance Study: The paper reports results across inference-class and training-class GPUs, explicitly disclosing where the system wins and where it trails established baselines.

Experimental Results

Correctness & Portability:

On the RTX 5090, A100, and H100, AMK reproduces HuggingFace greedy decode token-for-token and matches teacher-forced perplexity to $2.5 \times 10^{-7}$ on the SmolLM2-135M checkpoint.
The validator rejected all 6,091 unsafe schedules in the stress test, demonstrating a 0% false-accept rate.

Performance on Inference-Class GPUs (The Win):

RTX 5090 (Consumer): An auto-tuned int8 weight-only (W8A16) megakernel beats CUDA-graphed cuBLAS (bf16) by 1.19–1.23× at batch-1 decode. This is a precision-asymmetric comparison (int8 vs. bf16) where the win comes from streaming fewer weight bytes (0.61× the bytes), not a faster per-byte kernel.
Datacenter Inference Fleet: The int8 megakernel beats CUDA-graphed cuBLAS bf16 across the inference-class fleet:
- NVIDIA L4: Up to 1.33× faster (growing with model size).
- L40S: 1.25–1.27× faster.
- A10G: Up to 1.08× faster at scale.
Self-Improvement: The autoresearch loop autonomously improved the megakernel over its own default by 1.25–1.72× on the RTX 5090.

Performance on Training-Class GPUs (The Loss):

On high-bandwidth training-class GPUs (A100, H100), the auto-generated megakernel trails CUDA-graphed cuBLAS and vLLM.
- On H100, AMK was 1.65× slower than vLLM's default CUDA-graph path.
- On A100, the optimized bf16 GEMV reached only ~63% of the measured HBM peak (vs. cuBLAS's ~90%).
Diagnosis: The bottleneck is the per-tile cross-SM synchronization cost (a fixed cost the megakernel pays that cuBLAS does not). On high-bandwidth silicon, this synchronization overhead dominates, whereas on inference-class GPUs, the bandwidth savings from fusion and quantization amortize this cost.

Significance and Claims

The paper claims that AutoMegaKernel represents a shift from "raw speed" to systematic safety and portability in megakernel synthesis. Its primary contribution is a statically-checked compiler architecture that allows autonomous agents to safely explore the space of whole-model schedules.

Trust: It replaces the risk of hung GPUs with a static safety gate that rejects unsafe schedules before launch.
Generality: It achieves self-retargeting across architectures (sm_80/90/120) and model families without manual CUDA tuning.
Honesty: The authors explicitly frame the performance results as regime-dependent. The system wins on inference-class GPUs by fusing operations and leveraging quantization to reduce memory traffic, but it currently loses on training-class GPUs due to synchronization overheads that require coarser-grained scheduling (future work) rather than just better micro-kernels.

The paper concludes that while the current kernel quality does not yet match hand-optimized baselines on high-end training hardware, the statically-checked harness provides the necessary foundation to safely automate the future improvements (e.g., coarser synchronization, deeper pipelining) required to close that gap.