MegaBake GitHub

PyTorch compiler · MIT · Hopper

Your model.
One kernel.
Fully baked.

Compile supported Hugging Face models into one persistent CUDA kernel. Keep your model code. Keep torch.compile. Reduce kernel launch overhead.

Explore the benchmarks Get a megakernel
torch.compile 
compiled = torch.compile(
    model,
    backend=megabake.make_backend(),
)

Open source. MIT licensed. Built for NVIDIA Hopper.

THE WHOLE MODEL. ONE LAUNCH.
MODEL → MEGAKERNEL MegaBake mark: a layered cake with a flame on top CUDA / H200
Embedding → Transformer layers → Decode One persistent CUDA kernel

1 launch

per decode call. SmolLM-135M uses 842 launches with CUDA-graphed Inductor.

2.09×

decode speedup on SmolLM-135M versus torch.compile with CUDA graphs: 2.29 ms to 1.10 ms.

128 tokens

static KV-cache capacity for the published batch-1, bf16 decode comparison.

6 models

published decode packages on an NVIDIA H200 MIG 3g.71gb partition.

Measured 24 October 2026. The table and downloads use the same published reports from megabake/megakernels.

Measured 24 Oct 2026

Decode leaderboard

Six models, one GPU. Compare decode speed with CUDA-graphed Inductor, or open a breakdown for timings and model-specific notes.

H200 MIG · 60 SMs · bf16 · batch 1

Median complete call, including host and GPU work. Matches the published results.

MegaBake CUDA-graphed Inductor Speed differences within 10% are labeled close. Shared-GPU clock noise is 10–15%.
MegaBake batch-1 decode speed by model on an NVIDIA H200 60-SM slice.
Rank Hardware
1 SmolLM-135M
HuggingFaceTB · 1d46172
NVIDIA H200
MIG 3g.71gb · 60 SMs · bf16
912 tok/s
1.10 ms/token
2.09× faster109% faster
MegaBake
912
Inductor
436
2 Qwen2-0.5B
Qwen · main
NVIDIA H200
MIG 3g.71gb · 60 SMs · bf16
759 tok/s
1.32 ms/token
1.80× faster80% faster
MegaBake
759
Inductor
422
3 DeepSeek-V3 draft 0.6B
jukofyork · main
NVIDIA H200
MIG 3g.71gb · 60 SMs · bf16
747 tok/s
1.34 ms/token
1.76× faster76% faster
MegaBake
747
Inductor
425
4 SmolLM-1.7B
HuggingFaceTB · main
NVIDIA H200
MIG 3g.71gb · 60 SMs · bf16
419 tok/s
2.39 ms/token
1.30× faster30% faster
MegaBake
419
Inductor
323
5 Beam DFlash2 draft
reflection · main
NVIDIA H200
MIG 3g.71gb · 60 SMs · bf16
399 tok/s
2.51 ms/token
Within 10%Within 10%
MegaBake
399
Inductor
369
6 gemma-2b
google · main
NVIDIA H200
MIG 3g.71gb · 60 SMs · bf16
298 tok/s
3.35 ms/token
1.10× faster10% faster
MegaBake
298
Inductor
270

Source: megabake/megakernels, 24 October snapshot. Batch 1, bf16, one token per call, static 128-token KV cache, 40 samples. The baseline is a hand-captured CUDA graph of torch.compile(model, fullgraph=True), or Inductor’s reduce-overhead mode where faster. Speedup is baseline latency divided by MegaBake latency. Open each breakdown for its report. Long-context results use different cache sizes and are not mixed into this table.

Artifacts

Download a megakernel.

Download a published H200 decode package for each model: module.py, report.json, and source_map.json. Needs the MegaBake revision pinned in the report.

These are generated CuTe DSL sources, not standalone binaries or model weights. You need the model weights, a compatible Hopper GPU, and the matching MegaBake runtime. The package reports are the source of the leaderboard timings above. Upstream notes that the recorded compiler commit may have additional working-tree changes.

Source: megabake/megakernels. To get every kernel at once, download the full archive.

Integration & evaluation

A compiler backend.
A path to the Hub.

Keep the PyTorch and Hugging Face model interface. Evaluate the compiled workload, then distribute a verified kernel alongside its model revision.

Read the current compiler status

Supported today

Hopper / SM90, bf16, static shape buckets, and bounded contiguous decode KV state. Six published decoders in this benchmark; performance and numerical agreement vary by model. Broader cache formats and production serving need further work.

Reproduce before integrating

Start with the pinned revision in a package’s report.json. Check outputs across advancing decode calls and compare the complete call with CUDA-graphed Inductor on the same GPU, shapes, and context.

Benchmark harness ↗ Measurements & limitations ↗

Proposed Hub integration

Publish device bodies and revision-specific megakernels with explicit hardware, shape, runtime, and numerical requirements. This is an integration proposal; the downloadable packages currently live on GitHub.

Working today

  • Eight Hugging Face decoders compile into one persistent CUDA kernel for stateful decode. Prefill support varies by model.
  • Outputs are checked against PyTorch. SmolLM-135M and Qwen2-0.5B match eager bf16 across advancing decode calls.
  • Hopper, bf16. Static shapes, with megabake.bucketed covering a few prompt and batch lengths.

Next

  • Tile-level dataflow and cross-layer weight streaming.
  • Per-GPU autotuning, fused GEMM epilogues, paged KV cache, and batched serving.
  • Fast long-prompt prefill, multi-step and speculative decoding, FP8, Blackwell, and Ampere.

Why this matters

Compile through torch.compile.

A backend beside Inductor
MegaBake plugs in at the same handoff and keeps PyTorch’s guarantees about guards, aliasing, and state. The Dynamo/FX stack can schedule a whole model, not only emit one kernel per operator.
Support for new architectures
Models change every few weeks: new norms, RoPE variants, attention forms. MegaBake matches by semantics. Adding a model is an extension of the operator families.
Hardware as a body provider
The graph, the execution plan, and the device bodies are separate. Hopper is the body that runs today. Blackwell, Ampere, and other accelerators plug in behind the same compiler.
Compiled once, then distributed
A baked kernel is tied to one model revision, one shape bucket, and one GPU class, and it takes minutes to produce. The Hugging Face Hub already pins revisions and ships architecture-matched kernels. The proposal is to publish MegaBake’s device bodies on the Kernel Hub, and one verified megakernel per model revision next to the weights.