RightNow-AI/AutoMegaKernel

TL;DR: AutoMegaKernel (AMK) is an MIT-licensed harness that lets a coding agent (Claude Code / Codex) compile a HuggingFace Llama-family model into a single persistent CUDA megakernel — the whole decode forward pass in one launch — then autonomously search schedules and micro-kernels for latency, with every proposal gated by a static deadlock/race validator and full-model logit equivalence against eager PyTorch. Its headline result: an auto-found W8A16 int8 megakernel beats CUDA-graphed cuBLAS bf16 at batch-1 decode on inference-class GPUs (L4, L40S, A10G, RTX 5090), while the README openly states the equal-precision bf16 path still loses to cuBLAS.

Why this is interesting

Megakernels (whole model, one persistent launch, no per-op kernel launches or HBM round-trips between ops) are the right shape for latency-bound single-stream decode — voice, real-time, agentic loops — but hand-writing them is brutal and they're notoriously easy to deadlock. AMK's bet is architectural: make correctness a property of the structure rather than the generated code, so an LLM agent can safely search the schedule space unattended for hours. That's a plausible template for how agent-driven systems engineering gets done, beyond this one project.

How it works

Four layers with an explicit trust model:

• Layer 0, the VM: a hand-written, frozen persistent kernel — one threadblock per SM, a per-SM scheduler loop, page-based on-chip scratchpad, counter-based synchronization. Launched once per forward pass.
• Layer 1, instructions: ABI-conformant micro-kernels (GEMV/GEMM tiles, attention tile, RMSNorm, RoPE, SwiGLU, dequant), each unit-verified in isolation before it can enter a megakernel.
• Layer 2, the scheduler: HF model → graph IR → tiled task-DAG → instruction stream + page allocation. This is the agent's edit surface — a structured schedule object (tiling, fusion grouping, SM assignment, pipelining depth, page allocation), not CUDA code.
• Layer 3, dynamism (continuous batching, dynamic shapes, MoE): a roadmap placeholder.

Deadlock-freedom is by construction: producers only increment counters, consumers only wait on statically known thresholds, and the VM refuses any schedule that isn't a valid DAG — a bad agent proposal becomes a clean REJECTED rather than a hung GPU. They report zero false-accepts across 7,160 adversarial schedules. Two autoresearch loops run on top: one optimizing individual micro-kernels (their prior AutoKernel project), one optimizing the schedule IR — the new axis.

The results, read carefully

The int8 win is real but regime-specific: batch-1, position-0/low-context decode on inference-class GPUs, where decode is bandwidth-bound and int8 halves the weight bytes streamed per token.

The interesting diagnostic: the split isn't bandwidth-ordered (the 864 GB/s L40S wins by more than the 600 GB/s A10G) — it's a fixed per-tile cross-SM sync cost that larger GEMV-dominated models amortize on inference silicon but that training-class A100/H100 never overcome. And the README is unusually candid: on like-for-like bf16, AMK is ~1.24× slower than cuBLAS, sustaining ~51% of spec HBM bandwidth vs cuBLAS's ~90%. The int8 win comes from fewer bytes, not a better kernel. They even enforce this honesty mechanically — the bench refuses to emit a latency without a paired correctness verdict, and roofline distance is always reported.

What's actually there

Python package (hatchling, amk CLI), MIT. GPU megakernel verified on sm_80/sm_86/sm_89/sm_90/sm_120 from the same source; multi-token greedy decode with persistent KV cache matching eager token-for-token; a real checkpoint path (SmolLM2-135M reproducing HF generate exactly); 98 tests (78 CPU-only); MCP server + Claude Code skill/slash-commands for agent driving; measured datacenter numbers file-backed in the repo. A 10-minute unattended run self-improved its own schedule 1.47×. Paper on arXiv (2606.09682).

Gaps: Llama-family dense models only; no MoE, batching, or dynamic shapes (Layer 3 is a placeholder); wins limited to int8 batch-1 on inference GPUs; no claim against vLLM at throughput. The "data flywheel trains a learned prior" and "retargets to new silicon in days" claims are directional, not yet demonstrated at scale. 135 stars, one-month-old repo — early, but the engineering discipline (frozen VM, static validator, correctness-gated benchmarks) reads as genuine rather than demo-ware.

Try it

git clone https://github.com/RightNow-AI/AutoMegaKernel && cd AutoMegaKernel
uv sync

# CPU-only sanity check
uv run pytest && amk doctor && amk eval toy --device cpu

# With a CUDA GPU + nvcc
amk compile toy --gpu rtx5090 --regime single-stream
amk generate toy --gpu rtx5090 --prompt-ids "1,2,3" --max-tokens 32 --verify

# Real HF checkpoint end-to-end
uv run python examples/run_hf_model.py