FlashAttention

FlashAttention is an IO-aware exact attention implementation: it tiles $Q K^T$ into SRAM, uses online softmax with running max $m_i$ and normalization factor $l_i$ (recursively rescaled when merging blocks: $m_{new} = max(m, m_{local})$), never writes the $O(n^2)$ attention matrix back to HBM in the forward pass, and recomputes rather than stores in the backward pass; the motivation is the ~10-20x bandwidth gap between HBM and SRAM (A100: 40GB HBM vs ~108KB SRAM).

the de-facto attention implementation for LLM training and inference (PyTorch SDPA, NVIDIA FA2/FA3, DeepSeek/Meta stacks); gains grow with sequence length, and long-context training essentially requires it.

standard implementations shuttle the attention matrix between HBM and SRAM repeatedly, making IO the bottleneck; tiling + recomputation cuts HBM traffic from $O(n^2 d)$ to $O(\frac{n^2 d^2}{M})$ ($M$ = SRAM size), giving 2-4x speedup and $O(n)$ memory instead of $O(n^2)$, with results bit-identical to standard softmax (not an approximation).

Test Your Knowledge: Practice Questions for "FlashAttention"

Single choice pitfall questions with instant feedback and mistake tracking.