FlashAttention paper

The claim is about bytes, not flops

Attention has the same flop count however you compute it. The paper’s observation is that the naive implementation materialises the score matrix in HBM, reads it back for softmax, writes the probabilities, reads them again for the value product, and stores what the backward pass needs. On long sequences that traffic, not the arithmetic, sets the time. FlashAttention tiles Q, K and V into SRAM, keeps a running maximum and a running sum, and never writes the score matrix anywhere.

Exact, not approximate

Nothing is approximated. The online softmax changes the order in which terms are accumulated and rescales the accumulator when a new maximum arrives; in exact arithmetic the result equals the two-pass softmax. That is why the paper claims a wall-clock speedup with no quality caveat, unlike the approximate-attention line it compares against.

What it does not claim

It does not reduce the flop count, which stays quadratic. It barely helps at short sequence lengths, where there is little traffic to delete. The block-sparse extension is a separate contribution and can wait for a second pass.

What to read

The IO-complexity argument and the algorithm block. The complexity counting is the part a compiler engineer should be able to reproduce: count HBM reads and writes for naive and tiled attention at a given sequence length, head dimension and SRAM size.

Afterwards you should be able to say why the backward pass is cheaper in memory than storing the matrix, and which recomputation it pays for that.

Things that break

Believing the score matrix is never computed. It is computed, tile by tile, and discarded; only the storage is avoided. The tell is a claim that the method is linear in flops or that it approximates softmax. The second failure is applying it at short sequence lengths and reporting the regression as a bug.