What problem does FlashAttention solve, and is it an approximation?
FlashAttention makes exact softmax attention faster by being IO-aware: instead of writing the N-by-N score matrix to slow HBM and reading it back, it computes attention in tiles that fit in fast on-chip SRAM and carries a running softmax (max and denominator) from tile to tile. It is not an approximation — the result is identical to standard attention — it just moves far fewer bytes through memory, cutting memory from O(N^2) to O(N).
How to think about it
The key insight is that attention is memory-bound, not compute-bound on modern GPUs. The arithmetic in QKᵀ, softmax, and ×V is cheap; the cost is shuttling the N × N score matrix to and from HBM.
What standard attention does (and why it’s slow):
- Computes
S = QKᵀ(anN × Nmatrix) and writes it to HBM. - Reads it back to apply softmax, writes it again.
- Reads it once more to multiply by
V.
At N = 32k, that matrix is several GB per head and makes three round-trips through ~3 TB/s memory.
What FlashAttention changes:
- Loads blocks of
Q, K, Vinto SRAM (~20 MB, ~10× HBM bandwidth). - Computes attention one tile at a time, keeping a running max
mand denominatorℓ(the online softmax). - Rescales the accumulated output when a later tile has a larger score, so the final result equals full-row softmax exactly.
- Never writes the
N × Nmatrix to HBM; the backward pass recomputes tiles instead of storing them.
It’s the kernel behind PyTorch’s scaled_dot_product_attention and every serious inference stack — and the reason 128k-token context windows became practical.