AutoMegaKernel (AMK): Persistent CUDA Megakernel

• AutoMegaKernel (AMK) is a system that compiles HuggingFace Llama-family models into a single persistent cooperative CUDA megakernel, eliminating per-operator launch overhead.
• It features a layered architecture—combining a trusted VM, micro-kernels, and a static safety validator—to ensure deadlock- and race-free execution.
• AMK supports quantized execution and an agent-driven schedule search, achieving measurable performance gains in batch-1 autoregressive decoding.

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). Persistent threadblocks remain resident, synchronize in-kernel, and walk the model graph directly, reducing launch bubbles and avoiding inter-op HBM round-trips (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). All reported latency measurements are for batch size 1, decode position 0, empty KV cache, i.e. the weight-dominated first-token regime (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). The schedule IR is a typed SM-level task DAG synchronized by monotonic counters (Jaber et al., 8 Jun 2026).

The supported scope is explicit: bias-free projections, full rotary embeddings, SiLU-SwiGLU, RMSNorm, and GQA Llama-family models (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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” (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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-$$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$0 “which-producer” race (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). Validation throughput was about 5,150 schedules per second on CPU, and 24 accepted real lowerings matched eager bit-for-bit in the ReferenceVM (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). Program size scaled up to 3,410 IR tasks for TinyLlama-1.1B (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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

$$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$1

rounded in the abstract to $$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$2 (Jaber et al., 8 Jun 2026). The paper also reports a 64-token greedy decode that was byte-identical to model.generate for 64/64 tokens (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). Int4 is “correct-but-lossy” under naive round-to-nearest, with about 22% greedy-token agreement though still coherent text (Jaber et al., 8 Jun 2026). The blended total byte ratio for int4 is about $$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$3, i.e. a floor reduction of about

$$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$4

because only linear GEMV weights move from bf16 to int4, while tied embeddings, fp16 dequant scales, and non-GEMV buffers remain bf16 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). The loop is explicitly

$$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$5

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 (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). Second, there is a physical-floor guard: if a measured or predicted latency falls below the roofline floor $$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$6, the result is withheld as infeasible (Jaber et al., 8 Jun 2026). This caught a 1024-thread launch that a fallback cost model predicted at 87.6 $$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$7, below the 695 $$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$8 physical floor (Jaber et al., 8 Jun 2026).

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 $$\{\text{ptr}, \text{numel}, \text{rank}, \text{dtype}, \text{space}, \text{shape}[4], \text{stride}[4]\}.$$9 to 1991 $>8$0, a 1.25x drift-robust gain across 12 experiments with 6 kept (Jaber et al., 8 Jun 2026). A headless cold run improved from 2168 $>8$1 to 1261 $>8$2, a 1.72x gain by iteration 2, and a warm campaign seeded from prior runs started 1.41x faster than the cold start (Jaber et al., 8 Jun 2026). The system supports overnight unattended runs, basin-hopping after plateaus, resumability, crash recovery, and wake-up reports (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 $>8$3 the weight bytes (Jaber et al., 8 Jun 2026).

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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026). 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 (Jaber et al., 8 Jun 2026).

Relative to neighboring work, AMK sits in a distinct part of the megakernel design space. “Mirage Persistent Kernel” (Cheng et al., 22 Dec 2025) 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 (Cheng et al., 22 Dec 2025). “Ada-MK” (Dong et al., 12 May 2026) 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 (Dong et al., 12 May 2026). 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 (Jaber et al., 8 Jun 2026). 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.