How a training megakernel works
A training update makes a prediction and measures its error. It then changes the weights. The example below shows this with two multiplications. We follow that update through PyTorch and a Latch megakernel.
Use the controls to step through the diagrams. The diagrams explain execution. Measured timings are linked separately.
One training update
A model has adjustable numbers called weights. It uses these weights to turn an input into a prediction. A loss function measures the prediction error. Backpropagation computes how a small change in each weight would change the loss. The optimizer uses these gradients to update the weights.
The example has one input, two weights and one target.
- Start. The input is x = 2 and the correct answer is t = 1. The two weights begin at w₁ = 0.5 and w₂ = 1.5.
- Forward. Multiply: h = w₁ · x = 0.5 × 2 = 1. The value h is an intermediate: the backward pass will need it again, so it has to be kept somewhere.
- Forward. Multiply again to get the prediction: y = w₂ · h = 1.5 × 1 = 1.5.
- Loss. Measure the error: L = ½(y − t)² = ½ × 0.5² = 0.125. Training tries to make this number small.
- Backward. Start from the loss. How does it change if y changes? ∂L/∂y = y − t = 0.5.
- Backward. Gradient of the second weight: ∂L/∂w₂ = ∂L/∂y · h = 0.5 × 1 = 0.5. This reads h from the forward pass — the reason it was kept.
- Backward. Pass the gradient back through the multiplication: ∂L/∂h = ∂L/∂y · w₂ = 0.5 × 1.5 = 0.75.
- Backward. Gradient of the first weight: ∂L/∂w₁ = ∂L/∂h · x = 0.75 × 2 = 1.5. Now every weight has a gradient.
- Update. Move each weight against its gradient by a step size η = 0.1: w₁ ← 0.5 − 0.1 × 1.5 = 0.35 and w₂ ← 1.5 − 0.1 × 0.5 = 1.45.
- Check. Run the forward pass again with the new weights: h = 0.7, y = 1.015, L = 0.00011, down from 0.125. One update made the prediction much better.
From two weights to 200 million
TinyMoE applies the same process to matrices. It has 199,837,056 weights. Its forward pass includes attention and expert layers. In each expert layer a router sends each token to two of sixteen small feed-forward networks. AdamW updates the weights using running averages of the gradients and their squares. We time the complete training update at several batch sizes.
Each operation must wait for its inputs. The backward pass also needs values from the forward pass. These values must remain available until they are used. The GPU must therefore coordinate both execution and storage.
How GPUs usually run it
A GPU runs programs called kernels across its streaming multiprocessors (SMs). Our H100 has 132 SMs. PyTorch dispatches tensor operations to kernels. Compilation can combine several operations into one kernel. Intermediate results passed between kernels usually go through GPU memory.
There are three common ways to run the same sequence. Choose one to compare.
Eager PyTorch
Python issues each launch. The GPU can finish a small kernel before Python issues the next one. It then waits.
torch.compile
The compiler combines some operations. In this example it combines the loss with its gradient and fuses the optimizer operations. This reduces the number of launches. Results passed between kernels still go through GPU memory.
torch.compile with CUDA graphs
A CUDA graph records the launches once and replays them with one host call. The GPU no longer needs a call from Python for each launch. Each kernel still has to start and finish. Intermediate results still pass through GPU memory.
Our controls use compiled PyTorch. Some also use CUDA graphs.
TinyMoE’s published batch contains 16 microbatches of 8 × 1024 tokens. An optimizer update follows all 16 microbatches. In an earlier published-batch benchmark, compiled PyTorch launched 69,394 kernels per update. A reconstruction of the original Hugging Face Trainer setup launched 228,417. Latch launched one generated kernel per microbatch and 17 small scalar fills.
What Latch does
A training megakernel performs the forward pass, loss, backward pass and weight update in one persistent GPU program. That program stays resident across the SMs for the whole update. Latch generates it from model code.
1 · Describe the model and update
You supply the model as an nn.Module. You also supply an objective and an optimizer written as a pure tensor function. TinyMoE uses a short adapter over the public Hugging Face checkpoint. It preserves every parameter name and value. The adapter calls Latch’s attention and routing operations: bf16_sdpa.sdpa, token_topk, route_order and ragged_mm. Global-norm clipping and AdamW are tensor arithmetic passed to tensor_update.
We added TinyMoE’s routing and expert operations to the compiler and wrote the adapter. These are now included in the package. The build command generates the training program automatically.
2 · Capture records the whole update as tensor operations
Capture traces the model and update rule at the supplied tensor shapes. It records their operations and dependencies as a graph. This runs on CPU. TinyMoE’s published-batch program captured 32,010 operations.
captured = capture_module(model, objective, (tokens,), (labels, coefficients),
optimizer=tensor_update(update, init_state=initialize,
forward_state=True),
state=model_state)Abridged from megakernel_lab/tinymoe/build.py.
3 · PyTorch supplies the gradients
PyTorch autograd supplies the backward operations during capture. It applies the chain rule as in steps 5–8 of the example. Latch records these operations with the forward pass and optimizer. Latch provides derivative rules for its custom operations such as expert products and BF16 attention.
The graph records which forward values the backward pass needs. Expert gradients must also return to the tokens routed through each expert. The number of tokens per expert changes at every update.
4 · Scheduling and storage
Latch groups operations into phases across all 132 SMs. Operations within a phase can share intermediate values in registers or shared memory. A phase that needs results from other SMs must wait for them. Latch enforces this with a grid-wide barrier. Values passed between phases are stored in a planned GPU-memory arena.
The plan also fixes a numerical policy. TinyMoE's practical policy keeps FP32 master weights and Adam moments and uses BF16 matrix products and attention.
5 · Generation and compilation
Latch emits CUDA using reusable hand-written implementations of matrix products, attention, routing and reductions. It compiles this code with nvcc on a CPU-only machine. TinyMoE’s code is divided into 16 device units. These compile separately and link into one kernel.
TinyMoE runs as one cooperative kernel with 132 blocks. Each update also needs three scalar fills and five small copies. In the earlier 0.9.0rc1 measurements, a cold build at 2 × 512 tokens took 283 seconds on a dedicated 64-vCPU machine, including model loading, capture, generation and linking. The previous compiler took 544 seconds on the same machine. Both produced byte-identical output.
Use the generated program’s comparison and numerical-check commands before starting a new experiment.
Where the time goes
Combining operations removes kernel boundaries and some trips to memory. The arithmetic remains. Use the controls to compare the execution costs.
Separate kernels, replayed by a CUDA graph
One generated megakernel
| Cost | Separate kernels (CUDA graph) | One megakernel |
|---|---|---|
| Host launch | One graph replay; each kernel still starts and drains | One launch, plus small fills and copies |
| Dependencies | A kernel boundary between dependent steps | A grid-wide barrier where results cross SMs; every SM waits for the slowest |
| Intermediates | Written to GPU memory, read back by the next kernel | Fused steps keep them on chip; values crossing phases still use the arena in GPU memory |
| Weights and optimizer state | Read and written in GPU memory | Also read and written in GPU memory |
| Arithmetic | The full update | The same update, under the stated numerical policy |
| Operation quality | Vendor libraries such as cuBLAS and FlashAttention | Latch's own operation bodies, which can be slower |
The weights do not fit on chip
A megakernel can keep intermediate values in registers and shared memory. TinyMoE’s weights and AdamW state must stay in GPU memory. Together they occupy about twenty times the H100’s total on-chip storage. Each update must read and write them.
Operation size matters
Removing kernel boundaries helps most when each operation is small. Large matrix products depend more on the quality of the matrix and attention implementations. Grid barriers can also waste time: fast SMs must wait for slower ones.
Three kinds of workload
These three workloads illustrate where megakernels help and where they need more work.
Large regular matrices
Each large kernel can keep the SMs busy for a long time. Launch and boundary costs account for a small share of the update. Vendor matrix and attention libraries are highly tuned for this work.
Fusion can combine the elementwise operations between large matrix products. It can also avoid some writes and reads of intermediate values. The potential saving is small relative to the matrix work.
Latch’s matrix and attention implementations need to match cuBLAS and FlashAttention. Frequent barriers between small operations add further cost.
TinyMoE at its published batch (8 × 1024 tokens per microbatch): 977.2 ms versus 1,088.4 ms compiled, 1.11× faster Historical. Granite 3.0-1B-A400M at 4,096 tokens: 161.45 ms versus 97.89 ms, 1.65× slower Historical · not qualified. See also E2-TTT outer training, about 3× slower Research branch.
Many small routed experts
Conventional execution uses many small kernels to route tokens and compute expert outputs and gradients. Kernel boundaries can account for much of the time at small batches. Variable expert sizes also complicate CUDA-graph capture. The standard TinyMoE control uses compiled layers. The looped model has a separate control that captures forward and backward in a CUDA graph.
Latch combines routing, expert products and gradient scatter in one program. These operations no longer need separate launches or host calls.
Latch’s grouped products are slower than cuBLAS and CUTLASS grouped GEMMs. Deterministic gradient scatter is expensive when each token visits many experts. The overall advantage shrinks as batches grow.
TinyMoE-200M complete training update with the current compiler: 2 × 1024 tokens, 13.63–13.77 ms versus 37.50–38.42 ms for the conventional trunk-graph control, about 2.75× faster; 32 × 1024 tokens, 68.27–68.61 versus 114.02–114.77 ms, about 1.67× faster. These are same-session clocks at fresh, mid and late states, repeated on two H100s. Scope and numerical checks.
Sequential predict-and-learn
These learning steps must run in order. The conventional GradLev implementation uses about eighteen kernels per step even with CUDA graphs. The changing state passes through GPU memory between them.
A small changing state can be divided across SMs and kept in shared memory for the whole sequence. All learning steps then run in one launch.
Each step still needs grid barriers. The output head is too large to stay on chip and must be read from GPU memory at every token. Larger adapted blocks also exceed on-chip capacity.
GradLev on GPT-2 small, a complete block of 1,024 adaptation steps: 54.0 ms versus 189.5 ms for the strongest CUDA-graph control, 3.5× faster Measured on 0.9.0rc1; runnable example included. Online adaptation of TinyMoE's last two expert blocks: 1.26–1.67× slower than a compiled, CUDA-graph-replayed control Research branch.
Where the measurements are
The results page covers complete-update timing and training to a target quality. Get started with the included examples.