ForgeMegakernel: A General Framework for Efficient Auto-Regressive Model Decode Megakernels

• Current engines launch 280–448 kernels per token on an H100, and small RMSNorm, RoPE, or KV-store kernels account for one-fifth of latency in the 280-launch case.Kernel startup and retirement overhead makes decoding take 1.06–1.45× as long as a single persistent-kernel execution transferring the same bytes.Source: §1 · p. 1
• Handwritten megakernels reach 74–78% of H100 bandwidth but target one model, and adapting them to another requires weeks of specialist effort.Existing compiler- and agent-driven generation also face limited performance generalization or correctness guarantees.Source: §1 · pp. 1–2
• ForgeMegakernel guides generation with ten model-independent milestones and validates kernels using a standardized oracle that derives evaluation quantities and checks mid-states.The oracle checks byte counts, precision contracts, and a float64 error bound rather than relying on agent-defined checks.Source: §1 · p. 2
• Across 14 megakernels for eight models, MBU reached 50.5–85.9%, with geometric-mean speedups of 1.21× over SGLang 0.5.18 and 1.54× over a megakernel compiler.Inside SGLang, the megakernels decoded faster at comparable accuracy.Source: §1 · pp. 2–3
• The ablation study found that removing a milestone’s structural requirement while retaining its MBU floor reduced the final campaign result for every tested pair.The framework runs GPU evaluations in each generation iteration and uses measured latency, utilization, profiler counters, and error information to guide edits.Source: §1 · pp. 2–3

2.1 The decode step

• A decode step applies per-layer normalization, projections, RoPE, attention, and an MLP, with attention reading the key/value cache across prior positions.The step ends with a final RMSNorm and LM head.Source: §2.1 · p. 3

2.2 Model bandwidth utilization

• Model bandwidth utilization is B(b,S) divided by token time and peak HBM bandwidth, where B counts model-dependent memory reads for the configured batch and sequence length.The byte count includes gathered embedding rows rather than the full embedding table, which can inflate reported utilization by up to 1.22× for untied embeddings.Source: §2.2 · p. 3

2.3 Megakernels

• A megakernel keeps the decode step within one host launch, avoiding operator-boundary transitions; existing published designs are point designs settled for one model.Handwritten megakernels reach 74–78% of H100 bandwidth on the model for which they were written.Source: §2.3 · p. 3

3.1 Limited performance across cells

• The naive-prompt agent produced kernels slower than SGLang on all five cells, at 0.59–0.95 times its utilization and a geometric mean of 0.75.The prompt specified the model, cell, and single-launch requirement, but no structural properties.Source: §3.1 · p. 3
• A bandwidth target alone did not guide optimization: one campaign retained grid-barrier phases for 35 rounds and stopped at 34.2% MBU.That result exceeded SGLang’s 31% on the cell but remained below a handwritten megakernel’s 74.5%.Source: §3.1 · p. 4
• MPK was no faster than the agent across the same five cells, at 0.47–0.91 times SGLang’s utilization and a geometric mean of 0.65.Its utilization fell by about 44% when a ninth sequence entered a batch of eight.Source: §3.1 · p. 4

3.2 Unverifiable correctness of kernel fusing

• Output-only checks cannot localize defects, and accepted kernels lost 1.5, 2.5, and 1.0 GSM8K points on three tested model-and-cell settings.Equal-length prompts can conceal a bug that broadcasts one sequence’s decode position across a ragged batch.Source: §3.2 · p. 4
• Agent-defined checks also accepted kernels that lost 5.0, 2.5, and 1.5 GSM8K points on the three tested settings.Each session wrote its own tests and cleared an MBU bar, but there was no harness or gate to standardize the check.Source: §3.2 · p. 4

4 Design

• The milestone ladder specifies the structure to build, the oracle defines each verdict, and the generation framework drives both before deploying the kernel in a serving engine.Serving evaluation uses the engine’s scheduler, sampler, memory manager, and detokenizer on the same card, so per-token engine costs are included.Source: §4 · pp. 4–5, §4.3 · p. 8

4.1 The universal knowledge base and the milestones

• Ten ordered milestones move from functional checks through a single persistent launch to performance structure, with model-specific constants derived for each campaign.The performance core uses per-SM instruction streams, dependency counters instead of grid-wide barriers, and shared-memory buffer pooling; the ladder specifies obligations rather than a fixed kernel template.Source: §4.1 · pp. 5–6

4.2 The mid-state test oracle

• The mid-state oracle checks byte-accounted performance, numerical error, and arithmetic precision against criteria derived independently of the agent.It reads intermediate KV-cache states and logits against float64 references, while the protected evaluator fixes configuration-derived quantities before generation.Source: §4.2 · pp. 6–7

4.3 The end-to-end generation framework

• Each generation round uses a fresh coding agent to edit and measure the kernel on an H100, retaining or reverting changes from gate results.A separate review audits passing diffs before milestone advancement, and each ledger entry records the hypothesis, change, result, and decision.Source: §4.3 · p. 8

5.1 Setup

• Measurements use one H100 80GB GPU at 1980 MHz, with bf16 weights and KV cache, across 14 configurations from eight model families.The configurations span batch sizes 1–16 and contexts of 96–4096 tokens; baselines include SGLang 0.5.18, vLLM 0.27.1, llama.cpp, and MPK.Source: §5.1 · p. 8

5.2 Kernel experiment

• 50.5–85.9% MBU and geometric-mean speedups of 1.21× over SGLang and 1.54× over MPK across 14 configurations.Figure 3 compares the forged kernels with four baselines; the Llama-3.2-1B 𝑏1×96 kernel is excluded from aggregates.Source: §5.2.1 · p. 8, Figure 3 · p. 9
• The forged kernels’ mean log-probability error is 0.38–0.72 times SGLang’s across configurations, with a median ratio of 0.66.The error compares each decoder’s assigned log probability for the observed next token against a float64 forward pass.Source: §5.2.2 · p. 9
• SGLang incurs 301–653 μs of launch cost and 333–851 μs of synchronization waits per token; small-op launches add 225–455 μs.The megakernel launches once per step; its synchronization waits occur at dependency counters, and it performs fused small operations within the same kernel.Source: §5.2.3 · p. 10

5.3 Serving experiment

• Inside SGLang on GSM8K, 13 generated kernels achieve speedups of 1.03×–1.28×, with a 1.11× geometric mean and accuracy within 1.5 percentage points of baseline.The reported accuracy differences fall within run-to-run noise associated with ragged-batch scheduling.Source: §5.3 · pp. 10–11, Table 2 · p. 10
• On one GSM8K question, all four engines answer correctly, and the generated megakernel has the lowest measured per-token latency.Its measured speedups on that question are 1.04× over SGLang, 1.16× over vLLM, and 1.26× over llama.cpp.Source: Figure 6 · p. 10, §5.3 · p. 11
• Across sequence lengths 64–4096, the selected set of four Qwen3-8B kernels reaches 79.6–83.7% MBU, versus SGLang’s 72.8–74.1%.Each kernel has a compile-time batch size and context bound, so requests are routed among kernels by context length.Source: Figure 5 · p. 10, §5.3 · p. 11

Limitations

• The evaluation uses only an H100 80GB GPU; whether the milestones hold across other architectures and hardware configurations remains open.The authors identify shared-memory budget, core count, and launch model as dimensions for further experiments.Source: §6 · p. 11

Cost

• The median megakernel required 47 rounds, 14.7 agent-hours, 178 Mtok, and $177 to forge.The paper argues this cost is justified for large-scale or long-running serving deployments.Source: §6 · p. 11

Hand-written megakernels

• Handwritten megakernels report 78% of H100 bandwidth for a Llama-1B implementation and include FlashFormer, but target one model and configuration per kernel.The paper notes that adapting handwritten implementations to another model requires weeks of specialist effort.Source: §7 · p. 11

Compiled megakernels

• Megakernel compilers reuse programs across models, but MPK utilization declines as model size decreases and its implementation relies on model-dimension-derived constants.Other cited compiler approaches include Event Tensor, Ada-MK, and Fleet.Source: §7 · p. 11

Agent harnesses for kernels

• AutoMegaKernel reports 4–8% of the weight-bandwidth ceiling up to 1.1B parameters, while ForgeMegakernel reports 50.5–85.9% MBU from 0.6B to 13B inside a serving engine.Unlike the schedule-safety validator described for AutoMegaKernel, ForgeMegakernel’s oracle also checks numerical outputs, precision, and memory-traffic accounting.Source: §7 · p. 11

Conclusion

• Across 14 tasks on models spanning 0.6B–13B parameters, generated kernels reach 50.5–85.9% MBU, with 1.21× geometric-mean speedup over tuned SGLang and 1.54× over a megakernel compiler.Inside SGLang, the reported decoding speedup is 1.11× at comparable accuracy.Source: §8 · p. 12

A Milestone ablation

• All arms forged Qwen3-0.6B at b1×128 from the same seeded scaffold, agent, prompt, and thresholds, with a fixed budget of 50 rounds.The configuration’s MBU ceiling was 39.87%, and the retained thresholds were 28%, 35%, and 38%.Source: A Milestone ablation · p. 14
• 35.69% MBU for the control versus 26.54%, 29.60%, and 34.00% for arms without M6, M7, and M8, respectively; each ablated arm missed its retained floor.Figure 7 shows the four campaign arms.Source: Figure 7 · p. 14, A Milestone ablation · p. 14
• Without M6, the arm ended 1.5 points below its floor; the M7 arm was still climbing, while the M8 arm fell short of its floor.The authors characterize this as evidence for the direction, not the size, of each requirement: the study covers one campaign, model, and shape.Source: A Milestone ablation · p. 14

A second coding agent

• The second-agent run used codex exec on GPT-5.6-sol, with the seeded scaffold, rungs, thresholds, and prompt unchanged.It forged Qwen3-0.6B at b1×128.Source: B A second coding agent · p. 14
• At round 44 of 84, the campaign produced a working megakernel with 35.7% MBU, 89.5% of the cell’s ceiling, and 1.73× SGLang and 1.50× MPK on the same card.The decode step took 1009.6 μs end to end.Source: B A second coding agent · p. 14
• Default-harness runs reached 26.4% to 35.0% MBU for 194 to 378 Mtok, supporting ladder and oracle transfer to a different agent at comparable token cost.The reported comparison uses the strongest of the four default-harness runs for the same cell.Source: B A second coding agent · p. 14