Norm & activation backward
The forward norm (lesson 19's fused RMSNorm) was a clean per-row job. Its backward splits into two halves with opposite kernel personalities: a per-row gradient that recomputes the row statistic instead of storing it, and a per-feature weight gradient that is a reduction across every token in the batch. The second half is where the engineering is.
The question this lesson answers
RMSNorm forward was a streaming per-row kernel that fit on one screen. The backward of the same op is suddenly the kernel people get wrong. Why — and where does the difficulty actually live?
LayerNorm backward, derived
Forward (LayerNorm; RMSNorm drops the mean term): per row x ∈ ℝd, with μ = mean(x), σ = sqrt(var(x)+ε), x̂ = (x−μ)/σ, output y = γ·x̂ + β. Given dy, let dŷ = dy·γ. Then:
The dx formula has two reductions inside the row (mean(dŷ) and mean(dŷ·x̂)) — that is why dx couples every element of a row: a normalized output's gradient depends on the whole row, not just the matching element. But each row is independent, so dx parallelizes one-program-per-row exactly like the forward. The hard part is the other two lines.
Half one: dx recomputes the statistic, doesn't store it
To compute dx you need σ and x̂. You have two choices: save them from the forward, or recompute them from the saved x. The decision is lesson 25's save-vs-recompute trade, and recompute wins decisively here:
| Option | Extra HBM per row | Extra compute per row |
|---|---|---|
| Save x̂ (the full normalized row) | + d·2 bytes read (same size as x — doubles the norm's HBM) | none |
| Save σ only (1 scalar/row), recompute x̂ = (x−μ)/σ | + a few bytes | one elementwise pass over x (already loaded) |
| Save nothing extra; recompute μ, σ from x | 0 | one extra reduction over x (already loaded for dx anyway) |
Since the backward must read x regardless (it's in the dx formula), recomputing μ, σ is one extra in-register reduction over data already in SMEM — essentially free — and it saves a full d-sized HBM read that storing x̂ would cost. A norm is bandwidth-bound, so trading a tiny bit of on-chip compute for a halved HBM footprint is the right call. This is the recompute pattern from FlashAttention applied to the cheapest possible statistic.
Half two: dγ and dβ reduce across the whole batch
Here is the real kernel problem. dγ and dβ are per-feature (shape d), but each feature's gradient sums over every token: dγ[j] = Σrows dy[row,j]·x̂[row,j]. The reduction axis is the token axis — the very axis you parallelized dx over. So the blocks that compute dx in parallel each hold a partial contribution to the same d-vector, and those partials must combine. This is exactly the cross-block reduction tension from lesson 07, and there are three ways to resolve it:
| Strategy | How | Cost | Determinism |
|---|---|---|---|
| Atomic add to a [d] buffer | each block atomicAdds its partial dγ/dβ | contention on d slots when many blocks write; fp32 atomics | ❌ non-deterministic order → bitwise-variable |
| Partial buffers + reduce kernel | block g writes a row of a [num_groups, d] workspace; a second kernel sums the groups | extra HBM = num_groups·d, one extra pass | ✅ deterministic |
| Locked accumulation | per-feature lock (or grouped rows so each block owns disjoint features) | serialization on hot features | ✅ if order-fixed |
The production pattern (PyTorch, Apex, Liger) is the two-pass partial-buffer one: assign each program a contiguous group of rows, accumulate that group's partial dγ/dβ in registers, write one partial row per group to a small workspace, then a tiny second kernel reduces the num\_groups × d workspace to d. It is deterministic, and the workspace is small (a few hundred groups × d), so the extra HBM is negligible against the activation. Atomics are faster to write but you give up reproducibility — and lesson 30 explains why a training run that isn't bitwise-reproducible is a debugging nightmare.
Activation backward: fuse it into the epilogue
GELU/SiLU sit between matmuls in the MLP. Their backward is elementwise — dx = dy · f'(preact) — so the only question is whether dpreact round-trips through HBM. It shouldn't: fuse the activation derivative into the epilogue of the matmul-backward that feeds it, exactly as the forward fused activation into the matmul epilogue (lesson 19, Triton lesson 08). For SiLU, f'(x) = σ(x)·(1 + x·(1−σ(x))); you recompute σ(x) from the saved pre-activation rather than storing the derivative. One saved tensor (preact, or its sign for ReLU), one fused multiply, no extra HBM pass.
Worked number
RMSNorm backward, B·T = 16384 rows, d = 4096, bf16, on H100.
- dx: reads dy + x (2·16384·4096·2 = 0.27 GB), writes dx (0.13 GB), recomputes σ on chip. Bandwidth-bound: ~0.4 GB / 3.35 TB/s ≈ 120 µs.
- dγ/dβ: the same dy and x reads can be folded into the dx pass (one fused kernel does both halves), so the only extra cost is the tiny workspace (say 256 groups × 4096 × 4 B = 4 MB) and a sub-microsecond reduce kernel.
- If you had stored x̂ instead of recomputing: + 0.13 GB extra read ≈ +40 µs, a 33% slowdown on a bandwidth-bound kernel — for zero benefit.
Interactive · dγ/dβ reduction: atomics vs partial buffers
Set the batch rows and the group count (how many partial-buffer rows). The widget shows atomic contention (writers per feature slot) against the two-pass workspace cost, and flags the determinism trade.
What this sets up
You now have the two backward personalities — per-row independent (easy) and cross-batch reduction (the work) — and the recompute-the-stat reflex. The next lesson takes both to their hardest form: FlashAttention backward, where you recompute the entire N² score matrix, carry a per-row correction, and run the two-GEMM rule inside a streaming loop.