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).

Why is attention memory-bound? The bandwidth/capacity gap between HBM and SRAM?

How does online softmax merge block statistics (recursive formulas for running max m and normalizer l)?

Why recompute in the backward pass instead of storing? Memory from $O(n^2)$ to what?

What did FlashAttention-2 improve over v1 (dropping intermediate matrices, fewer non-matmul ops, ~2x)?

Why is FlashAttention exact? Essential difference from sparse/low-rank attention?