Abstract

LLM decoding on a GPU looks like one step, but it’s really hundreds of tiny kernel launches. The GPU spends much of its time between kernels rather than doing math.

MegaBake is an open-source (MIT) compiler that hooks into torch.compile. It takes a stock Hugging Face model and fuses its entire forward pass into one persistent CUDA kernel: embedding, every layer, the KV-cache updates, and the LM head.

This talk is about how that works, where it wins, and where the opportunities are.

Problem

Decoding is bound by memory bandwidth and launch overhead. Each operator launches as its own kernel, reads its inputs from global memory and writes its outputs back. Between kernels, the GPU drains and refills. CUDA graphs hide the launch cost but not the bubbles between kernels. Hand-written megakernels close that gap, but they’re written by hand for one model at a time.

Approach

• Hook into torch.compile and intercept the model just before PyTorch would turn it into hundreds of separate kernels.
• Replace each piece with a hand-tuned or generated GPU building block body.
• Stitch them all into a single kernel that stays resident on the GPU for the whole forward pass.
• Prove it’s correct, then race it against PyTorch’s best.

Takeaways

• Why kernel boundaries cost so much more than you’d expect in decode, and what a megakernel actually removes.
• How to plug a custom compiler into torch.compile without giving up the PyTorch and Hugging Face interface.
• Practical lessons on grid barriers, GEMMs for skinny matrices, and flash-decoding inside a persistent kernel.
• Where generated code still loses to hand-tuned code, and why.

Demo

Compile SmolLM-135M live, run generation through the megakernel, compare its tokens/s with graphed Inductor, and walk through the generated kernel source.