Source
Dao et al. (2022), LIT-074 — ARXIV-2205.14135.
Why recomputation is cheaper here than storing
Standard attention materialises the N×N score matrix, writes it to HBM, reads it back for the softmax, writes that, and reads it again for the backward pass. The arithmetic is trivial next to the traffic, so attention is memory-bound and the N² term is a bandwidth cost before it is a compute one.
FlashAttention keeps the block of scores in on-chip SRAM, consumes it immediately, and never writes it out. The backward pass then recomputes each block from Q, K and V — which is more FLOPs and less time, because the FLOPs were never the bottleneck. Memory falls from quadratic to linear in sequence length.
The accounting is LIT-074’s. Its Theorem 2 puts standard attention at
Θ(Nd + N²) HBM accesses against Θ(N²d²M⁻¹) for the tiled kernel, and with
head dimensions of 64–128 and roughly 100KB of SRAM, d² is many times smaller
than M. Its backward pass rebuilds each block from the running maximum and
normaliser saved in the forward pass rather than reading a stored matrix, and
the end-to-end gains it reports at unchanged FLOP count — 3× on GPT-2, 15% on
BERT-large — include that recomputation.
This is why the practice is not the usual activation-checkpointing trade. The ordinary version buys memory with compute at a real cost in step time; here the recomputation is close to free, because the thing it avoids is the expensive one.
It is still that trade’s attention case SOTA-249: the general rule — store activations at a subset of checkpoints and recompute the rest in the backward pass when activation memory is what binds — is the instruction, and this practice applies it to the layer where it pays best and, with blockwise kernels, costs least.
Where the recommendation came from
Child et al. (2019) recommends it three years earlier, in LIT-225 §5.4 — “we recompute the attention and feed-forward blocks during the backwards pass” — and for a different reason: ordinary activation checkpointing pays compute for memory, and attention at long sequence length is where that trade is best, because “memory usage is high for these layers relative to the cost of computing them”. Recomputation alone is what let that paper train networks of hundreds of layers at 16,384 context.
So the practice is older than its source, and what FlashAttention changed is
the price, not the instruction. That is the distinction ADR-029 adds
introduced_by: to hold: dropping Child et al. would not force a line of this
body to be rewritten, and would still leave the record dating a 2019
recommendation to 2022.
Condition
It holds while the kernel is memory-bound, which is where attention lives at the head dimensions transformers use. It also depends on the softmax being computable blockwise — the online, running-maximum formulation — so a variant that needs the whole row at once cannot be done this way.
The saved memory is what makes the long-context regime affordable at all; SOTA-086 is the constraint on how the blocks are sized.