FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Attention on GPUs is bound by memory traffic, not arithmetic. FlashAttention tiles the computation so the full attention matrix is never written to high-bandwidth memory, recomputing what it needs in the backward pass. The result is exact attention that is faster and uses linear memory in sequence length.
Key ideas
- Count memory reads and writes, not floating-point operations
- Online softmax lets attention be computed in blocks
- Recomputation can be cheaper than storing intermediates
Why read it now
It is the clearest example of treating a model as a systems problem. Longer contexts and cheaper training both trace back to this kernel.
Question to keep in mind
Which memory tier is the bottleneck on a GPU, and how does tiling change what has to cross it?