FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Reading guide

Problem and motivation

Transformers have quadratic time and memory complexity in sequence length $N$. Although approximate attention methods reduce theoretical FLOP counts, they often fail to achieve wall-clock speedups because modern GPU computations are primarily bottlenecked by memory transfers (IO) rather than arithmetic operations.

In standard attention, computing $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$, and $\mathbf{O}=\mathbf{PV}$ writes intermediate matrices of size $N\times N$ to slow High Bandwidth Memory (HBM) and reads them back multiple times. This incurs $\Theta (Nd+N^2)$ HBM accesses and $O(N^2)$ memory footprint, limiting practical context windows.

Key ideas

• IO-Awareness: Account explicitly for the GPU memory hierarchy: fast, small on-chip SRAM (e.g., 192 KB/SM with ~19 TB/s bandwidth on A100) versus large, slow HBM (40–80 GB with 1.5–2.0 TB/s bandwidth).
• Tiling (Block Softmax Computation): Split the input $\mathbf{Q},\mathbf{K},\mathbf{V}$ into blocks that fit entirely within SRAM. By tracking running softmax row-max and normalizer statistics, the softmax reduction is computed incrementally across blocks without materializing the full $N\times N$ matrix in HBM.
• Recomputation in Backward Pass: Avoid storing the intermediate $N\times N$ attention matrix for backpropagation. Instead, store only the output $\mathbf{O}$ and softmax statistics $(\ell ,m)$, recomputing $\mathbf{S}$ and $\mathbf{P}$ on-chip block-by-block in SRAM during the backward pass.
• Kernel Fusion: Fuse matrix multiplications, masking, softmax, and dropout into a single CUDA kernel to eliminate redundant HBM round-trips.

Method (with the key equations)

Given $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$, FlashAttention sets block sizes $B_c=\lceil M/(4d)\rceil$ and $B_r=\min (\lceil M/(4d)\rceil ,d)$, where $M$ is SRAM capacity.

Incremental Softmax Tracking

For a vector decomposed into sub-vectors $x^{(1)}$ and $x^{(2)}$, running statistics update via: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

For block $j$ of keys/values and block $i$ of queries:

1. Compute unnormalized scores on SRAM: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. Compute block local statistics: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$, and ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. Update running global normalizers: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. Update running output block ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

Complexity Analysis

• IO Complexity: Standard attention requires $\Theta (Nd+N^2)$ HBM accesses. FlashAttention requires $\Theta (N^2{d}^2{M}^{-1})$ HBM accesses. The authors prove that no exact attention algorithm can asymptotically achieve $o(N^2{d}^2{M}^{-1})$ accesses across all $M\in [d,Nd]$.
• FLOPs and Memory: FlashAttention performs $O(N^{2}d)$ FLOPs (with slightly more FLOPs in the backward pass due to recomputation) and requires only $O(N)$ auxiliary memory beyond inputs and outputs.

Extension: Block-Sparse FlashAttention

Restricting evaluation to non-zero blocks according to a predefined sparsity mask $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$ achieves an IO complexity of $\Theta (Nd+N^2{d}^2{M}^{-1}s)$, where $s$ is the fraction of non-zero blocks.

Results (with the most important numbers)

• Training Speedup:
  • BERT-large (seq. len 512): Outperforms the MLPerf 1.1 training speed record by 15% (17.4 min vs. 20.0 min on $8\times$ A100s).
  • GPT-2 (seq. len 1K): Achieves up to 3.5x speedup over HuggingFace and up to 1.8x speedup over Megatron-LM.
  • Long-Range Arena (LRA): Yields 2.4x speedup over standard attention; Block-sparse FlashAttention achieves 2.8x speedup.
• Longer Context Capabilities:
  • GPT-2 with 4K context: Trains faster than Megatron-LM at 1K context while improving validation perplexity by 0.7 points (17.5 vs. 18.2).
  • Long Document Classification: Scaling context from 512 to 16K yields a 4.3 micro-$F_1$ gain on MIMIC-III (52.8 to 57.1) and up to an 8.5 gain on ECtHR (72.2 to 80.7 at 8K).
  • Path-X and Path-256: FlashAttention is the first Transformer to solve Path-X (seq. len 16K) with 61.4% accuracy. Block-sparse FlashAttention achieves 63.1% accuracy on Path-256 (seq. len 64K).

Limitations and open questions

• CUDA Engineering Overhead: Implementing IO-aware kernels requires custom, low-level CUDA programming rather than standard PyTorch or compiler abstractions, making portability across hardware architectures challenging.
• Multi-GPU Attention: The optimal single-node IO schedule does not automatically extend to distributed multi-GPU attention across interconnects (NVLink, InfiniBand).
• Exact Attention Lower Bound: While optimal across general $M$, proving fine-grained parameterized complexity lower bounds in terms of $M$ remains an open problem.