FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Authors: Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
$\text{Memory IO Complexity: } O(N^2) \longrightarrow O\left(\frac{N^2 d}{M}\right) \text{ where } M \text{ is GPU SRAM capacity}$
Traditional attention writes the intermediate $N \times N$ attention matrix to GPU HBM, causing severe memory bandwidth bottlenecks. FlashAttention solves this via:
1. **Tiling & Block-Level Processing**: Divides Query, Key, and Value inputs into blocks that fit entirely within fast on-chip GPU SRAM (Static RAM).
2. **Online Softmax**: Computes softmax scaling factors dynamically across blocks using running maximums and sums, avoiding the need to store the full $N \times N$ attention matrix.
3. **Recomputation in Backward Pass**: Recomputes attention values on the fly during backpropagation from SRAM rather than saving huge activation matrices to HBM.
Frequently Asked Questions
What is GPU memory hierarchy and why does it matter?
GPU SRAM (on-chip cache) is ~10x faster than GPU HBM (main memory). Algorithms that minimize transfers between HBM and SRAM run dramatically faster even with identical FLOP counts.
Does FlashAttention change the numerical outputs of attention?
No. FlashAttention computes mathematically exact attention without any approximation, producing identical outputs to standard scaled dot-product attention within floating-point precision.