FlashAttention, read closely

The paper is usually summarised as "a fused attention kernel." That undersells it. The contribution is a decision about what not to store.

The claim

Standard attention materialises the full N by N score matrix in high bandwidth memory. For long sequences that matrix is the memory bottleneck and, less obviously, the time bottleneck, because writing it out and reading it back costs more than the arithmetic performed on it.

The move

Tile the computation so a block of the score matrix is produced, used, and discarded inside fast on-chip memory, never touching HBM. The softmax normalisation is the hard part, because it needs a global maximum and a global sum. The paper handles it with a running rescale, so each new tile corrects the partial result computed so far.

Why recomputation wins

The backward pass needs the scores again, and they were thrown away. So it recomputes them. On paper this is strictly more arithmetic. It is faster anyway, because the arithmetic is nearly free relative to the memory traffic it avoids.

This is the part worth carrying to other problems. When a kernel is memory bound, recomputing something is not a cost. It is a way of spending the resource you have spare.

What I had to work out myself

The tile size is not a free parameter. It is bounded by on-chip memory per streaming multiprocessor, and the paper's speedups assume you picked it against the hardware you are actually running. On a different GPU the same code can land well off the reported numbers, which is not a flaw in the paper but is easy to misread as one.