Why does FlashAttention-2's backward pass recompute attention, and what exactly does the forward pass save for it?

TL;DR: The forward saves Q, K, V, the output O and one FP32 log-sum-exp per query row, never P. For 32 heads at N = 8,192 that is 1 MiB of LSE per layer instead of 4 GiB of BF16 probabilities. Backward rebuilds each P tile as exp(S - LSE) on chip, and the softmax gradient needs only D = rowsum(dO * O), a vector computed before the main loop. Recomputation adds one GEMM to the four the gradient needs anyway, and it wins because it removes N-squared traffic from a step that was bound by it.

How to approach it

Write the gradient first, then ask which of its terms are N by N. Those are the ones the kernel must avoid storing. Then answer two questions separately: what lets P be rebuilt exactly, and what lets the softmax Jacobian be applied without a second N by N reduction. Mention why this comes up in inference loops: the same LSE is what split-KV decoding and cascade attention merge on, and post-training (RL rollouts and fine-tuning) runs these kernels on the same teams.

A strong answer

Forward, per head: $S = QK^\top/\sqrt d$, $P = \operatorname{softmax}(S)$ row-wise, $O = PV$. Given $dO$:

$$dV = P^\top dO, \quad dP = dO\,V^\top, \quad dS_{ij} = P_{ij}\big(dP_{ij} - D_i\big), \quad D_i = \textstyle\sum_j P_{ij}\, dP_{ij}$$

then $dQ = dS\,K/\sqrt d$ and $dK = dS^\top Q/\sqrt d$. Four of these are N by N: S, P, dP and dS. A naive autograd graph stores P (and often S) from the forward.

P for one layer, 32 heads, N = 8,192, BF16:
  8,192^2 x 2 B x 32 = 4 GiB per layer      (64 GiB per layer at N = 32,768)
LSE for the same layer, FP32:
  8,192 x 4 B x 32   = 1 MiB per layer

Two facts make storing P unnecessary.

P is a function of S and one scalar per row. The forward's online softmax ends with a running max $m_i$ and sum $\ell_i$. Saving $L_i = m_i + \log \ell_i$ is enough, because $P_{ij} = \exp(S_{ij} - L_i)$ exactly: no running max is needed in backward, since L is already the true row reference.

D does not need P. The row sum $\sum_j P_{ij} dP_{ij}$ looks like an N by N reduction, but

$$\textstyle\sum_j P_{ij}\, (dO_i \cdot v_j) = dO_i \cdot \sum_j P_{ij} v_j = dO_i \cdot O_i$$

so $D = \operatorname{rowsum}(dO \circ O)$, an N by d elementwise product. FA2 computes it in a small preprocessing kernel before the main backward loop.

import numpy as np

def attn_fwd(Q, K, V):
    S = Q @ K.T / np.sqrt(Q.shape[1])
    m = S.max(1, keepdims=True); P = np.exp(S - m); l = P.sum(1, keepdims=True)
    O = (P / l) @ V
    return O, (m + np.log(l))[:, 0]            # save O and LSE [N], never P

def attn_bwd(Q, K, V, O, lse, dO):
    c = 1 / np.sqrt(Q.shape[1])
    P = np.exp(Q @ K.T * c - lse[:, None])     # recomputed from Q, K and LSE
    D = (dO * O).sum(1)                        # rowsum(dO * O) = rowsum(P * dP)
    dV = P.T @ dO
    dP = dO @ V.T
    dS = P * (dP - D[:, None])                 # softmax Jacobian, one row at a time
    return dS @ K * c, dS.T @ Q * c, dV        # dQ, dK, dV

Against central finite differences on a 13 by 8 problem, all three gradients agree to within 1.7e-9, and the D identity holds to 7e-16. The real kernel tiles this: FA2's backward puts key blocks in the outer loop, so each thread block owns a K, V tile and accumulates dK and dV in registers, recomputes the P tile for every Q tile it visits, and adds its contribution to dQ with FP32 atomics, since many key blocks write the same dQ rows.

The cost of recomputation in FLOPs:

forward GEMMs:   QK^T, PV                          = 2
backward GEMMs:  QK^T (recompute), dV, dP, dQ, dK  = 5, of which 1 is recompute
backward / forward FLOPs                           = 2.5x

One extra GEMM out of five buys the removal of every N by N read and write. At N = 8,192 the stored-P version would move 4 GiB per layer each way through HBM; the recompute version stays compute-bound for the same reason the forward does.

The reversal: at short sequence lengths, a few hundred tokens, P is small, the backward is launch- and occupancy-limited rather than traffic-limited, and storing P can be cheaper. The case for recompute grows with N squared.

What interviewers probe next

"Why is LSE saved in FP32 and not BF16?" Its error goes straight into the exponent: an absolute error of 0.01 in L scales every recomputed probability in that row by about 1%. BF16 spacing near 10 is 0.0625.

"Where else is the LSE used in serving?" Merging partial attention results: split-KV decoding, cascade attention over a shared prefix, and context parallelism all combine (output, LSE) pairs with the same weighting.

"Why atomics for dQ rather than a second pass?" A second pass with Q blocks outermost would recompute P again. Atomics on an FP32 dQ buffer cost contention and non-determinism instead; deterministic modes trade speed for an ordered reduction.

How it's scored

Weak: says backward "recomputes to save memory" without naming what is saved, or claims recomputation makes backward cheaper in FLOPs.

Hire: writes the gradient, names O and the per-row LSE as the saved state, and explains why exp(S - LSE) reproduces P.

Strong Hire: derives D = rowsum(dO * O), counts the extra GEMM, sizes P against LSE in bytes, and describes the key-outer loop with atomic dQ.

Common mistakes

• Saving m and l separately and recomputing with a fresh running max, which is correct but wastes the point of LSE.
• Computing D from the recomputed P and dP inside the loop, which reintroduces an N-squared reduction per row.
• Forgetting the 1/sqrt(d) factor in dQ and dK, or applying it twice when it was folded into Q.
• Indexing the saved LSE by KV head in a GQA backward. It is per query head, and reading the wrong head's normalizer gives gradients that look plausible and are wrong.