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.