Latch — a compiler for training megakernels
Making unusual learning practical.
Suppose you’ve got an idea for a new learning algorithm.
The further it strays from familiar methods, the more of its implementation you’ll have to build yourself. Conventional methods have earned their place, and an entire ecosystem of kernels, libraries, and compilers has co-evolved with them through a decade of optimization. If your idea is radical enough, you have a systems project to finish before you can test it, let alone iterate on it.
Systems work holding research back, then enabling it, is nothing new:
- AlexNet relied on an efficient GPU implementation to make training practical.
- The Transformer removed recurrence and made sequence training much more parallel.
- FlashAttention made longer contexts practical by reducing how much data attention moves through GPU memory.
One exciting new primitive is the megakernel: fusing long sequences of GPU operations into a single kernel. This can reduce launch overhead and memory traffic while keeping the GPU busy between operations.
Hazy Research demonstrated this approach for whole-model inference. Cursor’s Mixture-of-Kittens applied it to mixture-of-experts training. Flapping Airplanes is pursuing megakernels as a route to more data-efficient learning, with early results earning them a $1B Series A.
Like most people reading this, I don’t have a billion dollars in venture funding. Maybe I’m also lazy. A training megakernel has to fit the forward pass, backward pass and optimizer into one program, and every change to the model means reworking it. That’s why I built Latch.
Introducing Latch
Latch is a compiler for training megakernels on H100. Give it supported PyTorch training code and it generates one GPU program for the forward pass, loss, backward pass and optimizer update, handling the dependencies, memory reuse and synchronization between them. When your experiment changes, you regenerate the kernel instead of rewriting it.
To demonstrate, I used Latch to train TinyMoE-200M. Generation and compilation took about five minutes on a 64-vCPU machine. The generated megakernel reached the same validation loss with 40% less training-update time than compiled PyTorch using the same recipe. At a small batch of 2 × 1024 tokens, complete training updates ran 2.75× faster.
What I’m most interested in are learning procedures that are impractical with today’s tools. In GradLev, a model updates its weights as it reads, one token at a time. Latch runs all 1,024 adaptation steps for GPT-2 small’s final MLP in a single kernel, keeping the changing weights in shared memory throughout. The complete block ran 3.5× faster than compiled PyTorch replaying a CUDA graph; both use ordinary PyTorch for the rest of GPT-2.
Latch has generated training programs for mixture-of-experts, state-space, diffusion and energy-based models, along with selective weight updates, repeated layers and token-by-token adaptation. Results.
It’s a research preview and a first step toward an open-source megakernel framework. Contributions are welcome. The open problems are a good place to start, and I’m curious to see what people explore with it.
Many thanks to Modal for generously supporting Latch with GPU credits!