FlashAttention backward

The flagship of the training Part. The forward (lesson 12) never materialized the N² attention matrix — so the backward has to recompute it, block by block, while carrying a single per-row correction scalar and running the two-GEMM rule inside a streaming loop. It composes every idea in Part IV: recompute-don't-store, the online softmax statistic, cross-block gradient accumulation, and fp32 reductions.

Builds on

Forward FlashAttention & online softmax (lesson 12); the two-GEMM rule (26); recompute-the-stat and cross-block accumulation (28). If lesson 12 was the forward flagship of the serving track, this is its mirror.

The question this lesson answers

FlashAttention forward earned its speedup by not storing the N² scores. That looks like it should make the backward impossible — gradients need the attention weights. How do you compute dQ, dK, dV when the thing you'd differentiate was thrown away?

What the forward left behind

Recall the forward output for one query row: O = softmax(QKᵀ/√d)·V, computed with the online recurrence so that the only things written to HBM are:

O — the output, shape N×d (you needed it anyway), and
L — the per-row log-sum-exp L_i = m_i + log ℓ_i, shape N (one scalar per query).

That's O(N·d) of saved state, not O(N²). The N² probability matrix P = softmax(S) is gone. The backward's whole trick is that L is exactly enough to regenerate any block of P on demand: P_ij = exp(S_ij − L_i), where S_ij = (Q_i·K_j)/√d is recomputed from Q and K. No global softmax pass needed — L already folded in the row max and normalizer.

The gradients, and the one scalar that makes them local

Given dO, the attention backward is:

dV = Pᵀ · dO dP = dO · Vᵀ dS = P ∘ (dP − D), where D_i = Σ_j dO_ij·O_ij (one scalar per row) dQ = (dS · K)/√d dK = (dSᵀ · Q)/√d

The pivotal term is D_i = rowsum(dO ∘ O). It comes from the softmax Jacobian: differentiating softmax gives dS = P∘(dP − (P·dP)·𝟙), and the inner product P·dP per row equals rowsum(dO∘O) — a quantity you can compute in one cheap pass over dO and O before the main loop. Precomputing D (shape N, like L) means each block of dS is computable locally from that block's P, dP, and the row scalar D_i — no second global reduction inside the loop. This is the same "hoist the per-row correction out of the loop" move as the online softmax, now on the backward side.

The loop structure: who accumulates cleanly, who needs atomics

Look at the reduction axes. dK_j and dV_j for a key/value block j sum over all query rows i. dQ_i for a query block i sums over all key blocks j. You can't make both the clean inner-loop accumulation at once. FlashAttention-2's backward chooses the outer loop over K/V blocks:

Outer loop j over K/V blocks; inner loop i over Q blocks. Inside, dK_j, dV_j accumulate in registers across all i and write once at the end of the j-iteration — clean, no atomics, fp32 accumulator.
dQ is the awkward one: each (i, j) contributes to dQ_i, so across outer iterations the same dQ_i gets many contributions. Two production choices: (a) atomicAdd the partial dQ_i into HBM each iteration (simple, non-deterministic), or (b) a separate dQ kernel that loops the other way (outer over Q, inner over K) so dQ accumulates cleanly too — the two-kernel split most libraries ship, trading a second recompute of S for determinism.

The trade in one line

FlashAttention backward spends extra FLOPs (recomputing S and P) to avoid O(N²) memory. Because attention is bandwidth/memory-bound at long context and the saved matrix grows quadratically while the recompute grows only with the work you were already doing, the trade is overwhelmingly worth it — and it's the only way attention training fits at long sequence at all.

Worked numbers: why O(N²) was never an option

One head, head-dim d = 128, bf16. The materialized probability matrix is N²·2 bytes per head per layer:

Seq length N	Materialized P (per head)	FlashAttn saved state O+L+D (per head)	Ratio
2,048	2048²·2 = 8.4 MB	2048·128·2 + 2048·4·2 ≈ 0.54 MB	~16×
8,192	8192²·2 = 134 MB	≈ 2.1 MB	~64×
32,768	32768²·2 = 2.1 GB	≈ 8.5 MB	~256×

Multiply by heads × layers and the materialized path is impossible past a few thousand tokens; the FlashAttention path scales linearly. The price is FLOPs: the backward's recompute of S makes attention backward ≈ 2.5× the forward attention FLOPs — paid gladly, because the alternative doesn't fit.

The traps

Trap	Symptom	Fix
Forgetting to precompute D	a global reduction sneaks into the inner loop; slow / wrong dS	compute D = rowsum(dO∘O) in a pass before the main loop
Recomputing P without L	numerically unstable exp; NaNs	P = exp(S − L); L already carries the row max
dQ accumulation race	non-deterministic grads or wrong dQ	atomics (accept non-determinism) or the separate dQ-major kernel
bf16 dK/dV/dQ accumulators	grad drift at long context	fp32 accumulators, downcast on store (lesson 26)
Causal mask handling in backward	gradient leaks across the causal boundary	skip fully-masked (i,j) blocks; mask the diagonal block

Interactive · the N² you don't pay, the FLOPs you do

Sweep sequence length. The widget contrasts the materialized P matrix (quadratic) with FlashAttention's saved state (linear), and shows the recompute FLOP overhead you pay for it — the recompute-for-memory trade made visual.

What this sets up

You've now built the gradient of the hardest forward kernel by composing every Part IV idea. All these kernels produce one thing: gradients. The last lesson consumes them — the optimizer kernel — and closes the training step back to lesson 25's memory peak, with mixed precision and recomputation as the final levers.