Part 5 · The Maths III: Training at Scale

FlashAttention: the same maths, different memory traffic

Attention is memory-bound, not compute-bound. The naive implementation writes the T×T score matrix to HBM, reads it back for the softmax, writes it again, and reads it once more to multiply by V. FlashAttention 2205.14135 never materializes it: it tiles Q, K and V into blocks that fit in on-chip SRAM and carries a running maximum and running sum so the softmax can be completed incrementally — the online-softmax trick of Milakov & Gimelshein 1805.02867.

The output is bit-for-bit the same attention. What changes is HBM traffic, and since SRAM bandwidth is roughly an order of magnitude higher than HBM, the wall-clock effect is large. It is an exact algorithm that is faster purely because of where the bytes live — the clearest example in the field of arithmetic intensity mattering more than FLOP count.