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

On GPUs, exact attention is often slow not because of arithmetic intensity but because of memory traffic. The textbook algorithm materializes the full $N \times N$ score matrix (and usually the attention probabilities) in HBM. For long sequences that matrix dominates both memory and wall-clock time — even when the FLOP count looks fine on paper.

FlashAttention’s punchline: treat attention as an IO-aware algorithm. Tile $Q$, $K$, and $V$ through fast on-chip SRAM, fuse the softmax with the matmuls, and never write the full attention matrix to HBM. The result is still exact attention (same math as softmax attention), just scheduled for the memory hierarchy.

What this paper explains

Standard attention for sequences of length $N$ and head dimension $d$ costs $O(N^{2}d)$ FLOPs but also $O(N^{2})$ HBM reads/writes if you store $S = QK^{\top}$ and $P = \mathrm{softmax}(S)$. HBM bandwidth is the scarce resource; SRAM is small but much faster.

Dao et al. show you can:

1. Load blocks of $K$ and $V$ (and streaming blocks of $Q$) into SRAM.
2. Compute partial attention scores and maintain running softmax statistics (online softmax / “safe” streaming normalization).
3. Accumulate the output block in SRAM and write only the final $O(Nd)$ outputs back to HBM.
4. In the backward pass, recompute attention instead of storing the huge intermediate — trading FLOPs for IO, which is a win when IO dominates.

Prior limits

• Materialized attention — simple to implement in frameworks, disastrous $O(N^{2})$ memory at long context.
• Approximate attention (sparse patterns, low-rank, kernelized) — can cut memory/FLOPs but change the function; quality and kernel fusion become research projects of their own.
• Kernel fusion without tiling — helps launch overhead, but still fails if the algorithm insists on writing $N \times N$ temps.

The mechanism

Think in blocks. Choose block sizes so a tile of $Q$, $K$, $V$ fits in SRAM. For each query block, iterate key/value blocks:

• Update local scores $S_{ij} = Q_i K_j^{\top} / \sqrt{d}$.
• Update running row-max and normalizing denominators so the final softmax matches the global one.
• Accumulate $O_i \leftarrow$ weighted $V_j$ contributions.

Because the algebra of softmax can be rewritten with running statistics, you never need all columns of a row resident at once — only the statistics that make the eventual normalization exact.

Algorithm / figure walkthrough

1. Partition $Q, K, V$ into SRAM-fitting blocks.
2. For each $Q$ block, initialize output accumulators and softmax stats to neutral values.
3. Stream $K, V$ blocks; update stats + output in SRAM.
4. Write the finished $O$ block to HBM; discard score tiles.
5. Backward: reload inputs, recompute block attention as needed, accumulate gradients.

The opening figure contrasts an HBM-resident $N\times N$ canvas with a small SRAM working set that slides over $Q/K/V$ tiles.

What to notice when reading

• The paper’s analysis is about HBM accesses, not just FLOPs — count bytes moved.
• Exactness comes from online softmax, not from dropping terms.
• Wall-clock gains depend on sequence length and hardware; the IO model predicts when tiling wins.

Results and evidence

The paper reports end-to-end speedups and large memory savings versus standard attention on long sequences (their benchmarks show up to roughly 2–4× faster attention and 10–20× lower memory for attention at long $N$, depending on setting). Lower memory enables longer context and larger batches; they demonstrate improved training throughput on language and long-document tasks relative to baseline attention implementations of the era.

Treat those ranges as the authors’ reported operating regime on the GPUs they used — not a universal constant for every model size.

Limitations

• Implementation is hardware-sensitive; block sizes must fit SRAM and warp constraints.
• Exact FlashAttention still has $O(N^{2}d)$ FLOPs — it does not asymptotically beat attention’s compute, only its IO/memory profile.
• Very short sequences may not benefit; overheads can dominate when $N$ is small.
• Later variants (FlashAttention-2, etc.) refine parallelism further; this paper is the IO-aware foundation.

How to read the paper

1. Abstract + §1 — HBM vs SRAM motivation.
2. §2 — background on attention and IO complexity.
3. §3 — tiling algorithm and online softmax.
4. §4–5 — runtime/memory experiments and ablations.
5. Appendix — numerical details / backward recomputation.

Knowledge check

Keep reading

1. Original paper (arXiv:2205.14135)
2. Related Fanout Daily: Attention Is All You Need, GQA, PagedAttention.

Sources

• Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022 — arXiv:2205.14135