AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis
2026-06-08 • Machine Learning
Machine LearningDistributed, Parallel, and Cluster ComputingPerformance
AI summaryⓘ
The authors created AutoMegaKernel (AMK), a system that turns a Llama-family HuggingFace model into one big CUDA program to run the whole model quickly without needing special hand-written code. They built a checker that stops unsafe schedules before they run, making the system reliable across thousands of tests. AMK works on multiple GPU architectures from one codebase and matches standard decoding results exactly on real models. The authors also show AMK can improve itself automatically and that its int8 megakernels outperform common GPU math libraries on many inference GPUs, though it is slower on very high-end training GPUs. This work focuses on the system design rather than just speed wins.
CUDA kernelLlama modelHuggingFacemegakernelstatic graph analysisdeadlock-freedomint8 quantizationinference optimizationGPU architectureperplexity
Authors
Jaber Jaber, Osama Jaber
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