---
title: 'AutoMegaKernel (AMK): Persistent CUDA Megakernel'
url: https://www.emergentmind.com/topics/automegakernel-amk
type: topic
---

# AutoMegaKernel (AMK): Persistent CUDA Megakernel

AutoMegaKernel (AMK) is a compiler-and-runtime system for turning an entire HuggingFace Llama-family forward decode step into one persistent cooperative CUDA megakernel, with a static safety gate designed so that an agent can search over schedules without being able to generate a kernel that deadlocks or races at runtime [2606.09682]. It 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, and presents its contribution explicitly as the system rather than raw speed [2606.09682]. The same source retargets `sm_80`/`sm_90`/`sm_120` from one codebase, auto-generates correct megakernels for 10 of 10 supported models, reproduces HuggingFace greedy decode token-for-token on a real SmolLM2-135M checkpoint, and combines a frozen validator, a hand-written trusted VM, a structured edit surface, and an agent harness for unattended schedule search [2606.09682].

## 1. Operating regime and performance rationale

AMK is specialized to batch-1 autoregressive decode, where one token requires streaming essentially the full model weight set once, so the arithmetic intensity is near 1 and the performance floor is dominated by memory bandwidth [2606.09682]. The paper writes that floor as
$$
t_{\min} = \text{weight\_bytes} / \text{HBM\_bandwidth}.
$$
In ordinary eager execution, or even CUDA-graph replay of eager execution, the model still runs as many separate kernels, so activations are written back to and re-read from HBM between operators, and host launch overhead is paid repeatedly [2606.09682].

A single persistent cooperative CUDA megakernel is therefore useful because it removes those operator boundaries: one launch is one forward pass is one decoded token [2606.09682]. Persistent threadblocks remain resident, synchronize in-kernel, and walk the model graph directly, reducing launch bubbles and avoiding inter-op HBM round-trips [2606.09682]. Execution uses one `cudaLaunchCooperativeKernel` with one threadblock per SM, and the KV cache persists in HBM across launches; thus each launch computes one token at a given decode position while preserving KV state across steps [2606.09682]. All reported latency measurements are for batch size 1, decode position 0, empty KV cache, i.e. the weight-dominated first-token regime [2606.09682].

## 2. System stack, compilation path, and scope

AMK is organized as a four-layer system in which correctness is treated as a property of the architecture, especially the validator and VM, not an emergent property of any one generated kernel instance [2606.09682].

| Layer | Role | Key content |
|---|---|---|
| Layer 0 | Trusted VM | Frozen, hand-written per-architecture persistent cooperative megakernel plus loader |
| Layer 1 | Instruction micro-kernels | GEMV, attention, RMSNorm, RoPE |
| Layer 2 | Scheduler/compiler | HuggingFace import, graph/task lowering, schedule validation |
| Layer 3 | Roadmap placeholder | Continuous batching, dynamic shapes, MoE |

The end-to-end pipeline starts from a HuggingFace Llama-family model, imports it within a specific scope, transforms it into a graph IR and then into a schedule IR, tiles the resulting DAG, assigns work across SM-resident queues, and lowers to a `MegakernelProgram` of tasks that map 1:1 to a fixed-size `amk_instruction_t` POD ABI [2606.09682]. The schedule IR is a typed SM-level task DAG synchronized by monotonic counters [2606.09682].

The supported scope is explicit: bias-free projections, full rotary embeddings, SiLU-SwiGLU, RMSNorm, and GQA Llama-family models [2606.09682]. Unsupported variants are mostly rejected at import; biased projections, linear RoPE scaling, and GELU activation are loudly rejected, while one honest blind spot is a Qwen2 variant with hardcoded `q/k/v` biases that passes the config-only importer check and then disagrees with eager execution [2606.09682]. This suggests that AMK is intentionally narrow at import time in order to keep later compilation and validation tractable.

The ABI is strict. Each task has an op code, up to 8 inputs, 4 outputs, 8 waits, one `out_counter`, an SM index, and a typed scalar parameter blob [2606.09682]. Buffers carry metadata
$$
\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.
$$
Inputs \(>8\), outputs \(>4\), waits \(>8\), or rank \(>4\) are rejected [2606.09682]. A Layer-1 instruction is pure compute: it may not touch counters, use undeclared buffers, or launch work; synchronization is owned entirely by the VM [2606.09682]. The agent-visible edit surface is not CUDA code but a structured `ScheduleConfig` with knobs for per-op tiling, fusion grouping, SM assignment strategy (`round_robin`, `load_balance`, or explicit), pipelining depth, page allocation (`linear`, `graph_color`, or `none`), threads per block, and shared memory bytes per block [2606.09682].

## 3. Schedule IR, persistent execution, and static safety gate

AMK’s synchronization model is deliberately simple. Cross-task communication is only through monotonic `uint32` counters [2606.09682]. When a task finishes, it performs a device-scope release fence ordering all output stores, then increments exactly one output counter by 1; semantically, that means “all my outputs are written and visible” [2606.09682]. A consumer task waits on a set of statically known \((\text{counter}, \text{threshold})\) pairs using an acquire-load spin loop, with exponential backoff and an abort-flag poll for WDDM/TDR escape [2606.09682].

At runtime, all blocks enter the persistent scheduler loop; two grid-wide barriers bracket execution so host-zeroed counters are visible before spinning and all stores are visible after completion [2606.09682]. Each SM queue is serialized in a global topological order, and the scheduler loop executes: optional prefetch of the next GEMV tile if pipelining is enabled, wait on all dependencies, cooperative dispatch of the Layer-1 compute kernel within the block, signal of the output counter, and continuation to the next instruction [2606.09682].

The validator is the core safety mechanism. It is a frozen static checker in `schedule/ir.py` that returns a `ValidationResult` and never throws; if validation fails, launch must not occur [2606.09682]. The authors are careful to state that this is not a mechanized proof, but a set of static graph checks over a trusted, hand-written base: the validator itself and the VM [2606.09682]. For deadlock-freedom, it enforces wait satisfiability
$$
1 \le t \le \#\text{producers},
$$
checks DAG acyclicity of the producer-to-consumer graph, and checks per-SM queue order so an SM cannot stall waiting on a counter that only a later task in the same queue would increment [2606.09682]. For race-freedom, it imposes a shared-counter all-join rule: if a counter has multiple producers, then any wait on that counter must use the full threshold equal to the producer count; a partial wait
$$
1 < t < \#\text{producers}
$$
is rejected as a first-\(k\)-of-\(N\) “which-producer” race [2606.09682]. It also computes transitive happens-before provenance for data hazards, adds an extra rule for same-pass `KV_CACHE` reads, and checks that any `IO_OUTPUT` buffer is produced by some task [2606.09682].

The adversarial evaluation is unusually detailed. Across 7,160 schedules, including 360 real lowerings, 2,800 single-injection mutants spanning 8 unsafe classes, and 4,000 random DAGs, an independent structural-plus-dynamic oracle labeled 6,091 as unsafe; the validator rejected all 6,091 unsafe schedules, for zero false accepts, and accepted all 360 real lowerings [2606.09682]. Validation throughput was about 5,150 schedules per second on CPU, and 24 accepted real lowerings matched eager bit-for-bit in the ReferenceVM [2606.09682]. One important controversy is addressed directly by the paper itself: AMK’s “correctness-by-construction” applies at the schedule level under the trusted-base assumption, not as a machine-checked proof of the entire CUDA program or every micro-kernel [2606.09682].

## 4. Correctness, retargeting, and quantized execution

AMK’s correctness story is anchored by a CPU-side `ReferenceVM`, exercised together with the importer, lowering path, validator, and CUDA path [2606.09682]. For generation coverage, AMK auto-generated correct megakernels for all 10 supported models tested: three real checkpoints—SmolLM2-135M, SmolLM2-360M, and TinyLlama-1.1B-Chat—and several from-config Llama shapes from about 40M to 618M parameters [2606.09682]. Program size scaled up to 3,410 IR tasks for TinyLlama-1.1B [2606.09682].

Retargeting is demonstrated on three architecture families. The same source generated and ran correct megakernels on RTX 5090 Laptop (`sm_120`), A100-SXM4-40GB (`sm_80`), and H100-80GB (`sm_90`), with `gencode` derived from the live device [2606.09682]. On the real SmolLM2-135M checkpoint, AMK reproduced HuggingFace greedy decoding token-for-token and matched teacher-forced perplexity over 187 real next-token predictions essentially exactly: 14.948473 versus 14.948473, with absolute gap
$$
2.45\times 10^{-7},
$$
rounded in the abstract to \(2.5\times 10^{-7}\) [2606.09682]. The paper also reports a 64-token greedy decode that was byte-identical to `model.generate` for 64/64 tokens [2606.09682].

Quantized generation is supported through the same compiler path. AMK auto-generates int8 and int4 weight-only quantized megakernels, folding dequantization into GEMV rather than requiring separate hand-coded kernels per precision [2606.09682]. The int8 path is described as greedy-lossless: on the real SmolLM2-135M checkpoint, it achieved 100% token agreement over 32 tokens relative to fp16 [2606.09682]. Int4 is “correct-but-lossy” under naive round-to-nearest, with about 22% greedy-token agreement though still coherent text [2606.09682]. The blended total byte ratio for int4 is about \(0.41\times\), i.e. a floor reduction of about
$$
2.42\times,
$$
because only linear GEMV weights move from bf16 to int4, while tied embeddings, fp16 dequant scales, and non-GEMV buffers remain bf16 [2606.09682].

## 5. Agent harness and unattended autoresearch

The agent harness is the outer propose/evaluate/keep-revert loop that lets a coding agent manipulate a constrained configuration surface—schedule knobs and kernel knobs—while the frozen compiler, validator, correctness oracle, and benchmark machinery reject bad candidates and preserve only measured improvements [2606.09682]. The loop is explicitly
$$
\text{propose} \to \text{validate} \to \text{correctness-gate} \to \text{measure} \to \text{keep/revert}.
$$
At Layer 2, the candidate consists of a `ScheduleConfig` plus GEMV build knobs such as `cols_per_warp`, tile width (`N_tile` in the experiments), threads, and `cp.async` depth [2606.09682].

Two robustness mechanisms are central. First, keep/revert uses interleaved back-to-back A/B remeasurement of candidate versus incumbent, so improvements caused merely by clock drift or thermal ramp are not accepted [2606.09682]. The paper notes on the unpinned laptop GPU that the same “best” config could show single-shot speedups from 1.08x to 1.39x across rounds, while the drift-robust accepted median was 1.25x [2606.09682]. Second, there is a physical-floor guard: if a measured or predicted latency falls below the roofline floor \(\text{weight\_bytes}/\text{HBM bandwidth}\), the result is withheld as infeasible [2606.09682]. This caught a 1024-thread launch that a fallback cost model predicted at 87.6 \(\mu s\), below the 695 \(\mu s\) physical floor [2606.09682].

The reported autoresearch gains are intentionally self-relative. On a 622 MB “small” model on RTX 5090, an agent-in-the-loop run improved from 2514 \(\mu s\) to 1991 \(\mu s\), a 1.25x drift-robust gain across 12 experiments with 6 kept [2606.09682]. A headless cold run improved from 2168 \(\mu s\) to 1261 \(\mu s\), a 1.72x gain by iteration 2, and a warm campaign seeded from prior runs started 1.41x faster than the cold start [2606.09682]. The system supports overnight unattended runs, basin-hopping after plateaus, resumability, crash recovery, and wake-up reports [2606.09682]. This suggests that AMK’s agent layer is closer to a constrained synthesis harness than to free-form code generation.

## 6. Performance envelope, limitations, and placement in the mega-kernel landscape

The main positive performance result is regime-specific. A search-found int8 weight-only megakernel (`W8A16`) beats CUDA-graphed cuBLAS bf16 for batch-1 decode on inference-class GPUs [2606.09682]. On L4, the median speedup grows with model size: 1.18x at 1.3B, 1.25x at 2.7B, 1.32x at 3.5B, and 1.33x at 4B [2606.09682]. On L40S, AMK wins by 1.25x at 4B and 1.27x at 6.7B; on A10G, it reaches parity at 2.7B and then 1.04x at 3.5B and 1.08x at 4B; on the consumer RTX 5090, AMK int8 beats cuBLAS bf16 by 1.19x on the 4-layer 623 MB model and 1.23x on the 8-layer 984 MB model [2606.09682]. The paper is explicit that this is not because AMK’s kernel is more efficient per byte: the equal-precision control shows AMK bf16 trails cuBLAS bf16 on RTX 5090 by 0.76–0.88x, so the int8 win comes from streaming fewer bytes, roughly \(0.61\times\) the weight bytes [2606.09682].

The losses are equally explicit. On training-class A100 and H100, the auto-generated megakernel trails cuBLAS, both in bf16 and in the self-tuned int8 regime [2606.09682]. The paper localizes the current boundary to per-tile cross-SM synchronization cost rather than simply memory bandwidth, and says direct A/B tests of `cp.async` int8 staging rings and split-KV attention both regressed [2606.09682]. The strongest caveats are also stated plainly: the cuBLAS comparison is precision-asymmetric (`W8A16` versus bf16), decode timings are batch-1 at position 0 with empty KV cache, the safety validator is not a formal proof system, the largest real checkpoint tested is TinyLlama-1.1B, and long-context decode was not measured [2606.09682].

Relative to neighboring work, AMK sits in a distinct part of the megakernel design space. “Mirage Persistent Kernel” [2512.22219] introduces an SM-level graph representation and an in-kernel parallel runtime and reports up to 1.7x over kernel-per-operator LLM serving systems [2512.22219]. “Ada-MK” [2605.11581] combines a three-dimensional shared-memory constraint model with MLIR-based fine-grained DAG offline search and improves single-batch throughput by up to 23.6% over vanilla TensorRT-LLM and 50.2% over vLLM on an NVIDIA L20 [2605.11581]. AMK’s own paper positions itself against MPK and hand-built Llama megakernels by emphasizing automatic generation with zero per-model CUDA plus static schedule checking and retargetability from one codebase [2606.09682]. This suggests that AMK’s defining contribution is the combination of self-retargeting whole-model megakernel generation, a frozen static deadlock/race gate, and an agent-safe edit surface rather than a claim of universal speed leadership.

Source: https://www.emergentmind.com/topics/automegakernel-amk