Matmul backward — the two-GEMM rule

Every learnable layer is a matmul, so every backward pass is built from matmul backward. One forward GEMM becomes two backward GEMMs of different shapes, with transposed operands and a reduction over the batch dimension. Get this kernel's layout and accumulation right and the rest of the backward is variations on it.

The question this lesson answers

Forward is one line: Y = X·W. You profile a training step and see the GEMM time roughly triple versus inference. Where did two extra matmuls come from, what shapes are they, and which one is the troublemaker?

Deriving the two GEMMs

Take Y = X·W with X ∈ ℝ^T×K (T = batch·tokens, K = in-features), W ∈ ℝ^K×N, Y ∈ ℝ^T×N. Given the upstream gradient dY = ∂L/∂Y ∈ ℝ^T×N, the chain rule gives:

Three outputs, three kernels. Two of them are full GEMMs the same size as the forward; one is a cheap reduction. Notice what changed:

Same FLOP count three times (2·T·K·N) — that's the 2× of lesson 25. But the two backward GEMMs are not interchangeable, and the difference is the contraction dimension.

Why dW is the special one: it contracts over the batch

The contraction (reduction) dimension is the one you sum over. Forward and dX contract over a feature dimension (K or N) — typically a few thousand. But dW = Xᵀ·dY contracts over T = batch·tokens, which at training scale is tens of thousands to millions. Two consequences follow, and the first is unconditional:

1. fp32 accumulation is mandatory. Each entry of dW sums T products — "this weight's contribution from every token in the batch." Summing tens of thousands of bf16 products in bf16 loses the low bits exactly where it hurts — the same trap as lesson 19's sum-of-squares. Accumulate dW (and db) in fp32 even when X, dY, and the stored gradient are bf16. This bites at every weight shape, because T is always large.
2. When the weight is small, dW under-fills the GPU and wants split-K. A GEMM is tiled over its output, and dW's output is the weight, K×N. For a fat MLP weight (4096×14336 → thousands of 128×128 tiles) that fills the SMs fine. But for a small output projection — a gate, a low-rank head, an attention out-proj at modest d — K×N tiles into fewer blocks than there are SMs, and the GPU starves while a huge T-reduction runs inside each. The fix is split-K: partition the T-contraction across blocks, each computing a partial dW, then sum the partials. That sum is a cross-block reduction → either atomics (fast, non-deterministic) or a second reduction kernel over a workspace buffer (deterministic, one extra pass). Lesson 30 explains why determinism matters for training. Note that split-K helps only because the output is small — growing T makes the reduction longer, not the occupancy worse.

The transpose: don't physically transpose, re-stride

Both backward GEMMs name a transposed operand (Wᵀ, Xᵀ). The naive move — materialize the transpose in HBM first — costs a full read+write of the tensor and a non-coalesced access pattern (lesson 04). Nobody does that. Instead the GEMM kernel reads the operand with swapped strides: a "transpose" is just telling the tile loader that rows and columns trade roles. cuBLAS exposes this as the transA/transB op flags; in Triton it's which stride you multiply tl.arange by when you build the pointer block. The data never moves.

The catch is on tensor cores (lesson 09): the mma instruction wants its fragments in a specific layout, so one transpose orientation lowers to a clean mma and the other may need an in-SMEM transpose or a layout conversion. This is why a backward GEMM can be measurably slower than the identically-sized forward GEMM even though the FLOPs match — the profiler shows it as extra SMEM traffic or a layout-conversion epilogue, not extra math.

Accumulate into the existing gradient (+=)

Weight gradients are not assigned, they are accumulated: W.grad += dW. Two reasons, both structural:

• Gradient accumulation across micro-batches. To simulate a large batch on limited memory you run several micro-batch forward/backwards and sum their gradients before one optimizer step. The dW kernel must add into the buffer, not overwrite it.
• Weight sharing. A tied embedding/unembedding, or any weight reused at multiple points in the graph, receives a gradient contribution from each use; they sum.

So the dW epilogue is a read-modify-write of the fp32 grad buffer, and the split-K partials must also add (not assign). This is why the grad buffer is commonly fp32 even in a bf16 model — repeated accumulation into bf16 would drift.

Worked example: one MLP layer of a 7B model

Take the up-projection of a transformer MLP: K = 4096 (d_model), N = 14336 (d_ff), batch·seq T = 8192, bf16.

Three ~0.96 TFLOP GEMMs ≈ 2.9 TFLOP per layer per step. At ~490 TFLOP/s sustained that is ~5.9 ms/layer of GEMM (versus ~2.0 ms for the forward GEMM alone) — the backward's two extra GEMMs are the reason the per-layer GEMM cost roughly triples from inference to training. At this fat MLP shape the dW output (4096×14336) tiles into thousands of blocks and fills the GPU fine — split-K isn't needed. It's the small output projections (the widget below lets you shrink the weight to find them) where dW's K×N output starves the SMs and split-K over the T-contraction pays off. The fp32 accumulation, by contrast, you need at every shape.

The traps, collected

Interactive · the three GEMMs and where dW chokes

Set the layer shape and batch. The widget draws forward, dX, and dW with their FLOPs and shows dW's output-tile count — the measure of how badly it under-fills the GPU before split-K. Watch dW's occupancy collapse as the weight shrinks and the batch grows.

What this sets up

You can now read any backward GEMM: identify the contraction dimension, spot the transpose, and predict whether it needs split-K and fp32 accumulation. Next we apply this to the single biggest tensor in a language model — the logits — where the forward "matmul into a loss" and its backward must be fused or the [B·T, V] softmax and its gradient blow your memory budget.