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

Key Idea

Standard attention computes the full N×N attention matrix, which is slow not because of compute (FLOPs) but because of memory reads/writes (IO) between GPU SRAM and HBM. FlashAttention uses tiling and kernel fusion to compute exact attention without ever materializing the full attention matrix in HBM.

The IO Bottleneck

• GPU compute has grown much faster than memory bandwidth
• Standard attention writes O(N²) intermediate values to slow HBM, then reads them back, this IO dominates runtime
• The attention matrix itself (N×N) is the bottleneck for both memory and speed

How FlashAttention Works

1. Tiling: split Q, K, V into blocks that fit in SRAM (fast on-chip memory)
2. Compute attention block by block: process one tile of Q against all tiles of K, V
3. Online softmax: compute softmax incrementally across blocks without needing the full row
4. Never materialize N×N: the full attention matrix never exists in HBM
5. Kernel fusion: combine all operations (matmul, softmax, dropout, matmul) into a single GPU kernel

Why It Matters

• 2-4x wall-clock speedup over standard attention implementations
• 5-20x memory reduction: linear memory in sequence length instead of quadratic
• Enables much longer context windows, training with 16K, 64K, 128K+ tokens became practical
• Now the default attention implementation in PyTorch, HuggingFace, and all major LLM training frameworks
• FlashAttention-2 and FlashAttention-3 further improved performance with better parallelism

Key Takeaways for Interviews

• The bottleneck is IO (memory bandwidth), not compute (FLOPs), this is the key insight
• FlashAttention computes exact attention (not approximate), no quality loss
• Tiling + online softmax + kernel fusion = never materialize the N×N attention matrix
• Enabled the long-context revolution: GPT-4's 128K context, Claude's 200K context, Gemini's 1M+ context