FlashAttention

Standard attention writes the whole score matrix to memory, then reads it back to apply the softmax, then reads it again for the weighted sum. FlashAttention never writes it at all: it works through the sequence in tiles small enough to stay in fast on-chip memory, accumulating the result as it goes. Identical answer, a fraction of the memory traffic.

Tile the computation and use an online softmax that can be updated as new blocks arrive, rescaling the running total rather than needing all scores first. The gain is in memory movement, not arithmetic — attention was bandwidth-bound, not compute-bound, which is why the speedup arrives with no change to the maths.

Online softmax keeps a running maximum m and sum ℓ, rescaling the accumulator by exp(m sub _old  − m sub _new ) as each block is processed. Memory traffic drops from Θ(n²) to Θ(n²d²/M) for on-chip memory of size M, while the output is exact rather than approximate.