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

Dao et al. (2022) — ARXIV-2205.14135

Key takeaways

• The bottleneck is HBM traffic, not FLOPs. Approximate-attention methods reduced compute complexity and often failed to produce a wall-clock speedup, because the cost model they optimised was the wrong one
• Tiling to on-chip SRAM, with an online softmax carrying a running max and normaliser, computes exact attention without materialising the N × N matrix
• Theorem 1: O(N²d) FLOPs — unchanged — and O(N) additional memory
• Theorem 2: standard attention needs Θ(Nd + N²) HBM accesses; FlashAttention needs Θ(N²d²M⁻¹), for d ≤ M ≤ Nd. With d ∈ [64,128] and M ≈ 100KB, d² ≪ M
• Proposition 3: no algorithm computes exact attention in o(N²d²M⁻¹) HBM accesses across all M ∈ [d, Nd]. It is IO-optimal, not merely faster
• Block sizes are derived, not tuned: B_c = ⌈M/4d⌉, B_r = min(⌈M/4d⌉, d)
• Backward recomputation rather than storing the attention matrix — trading arithmetic for memory traffic, which is the right way round given Theorem 2
• 15% over the MLPerf 1.1 BERT-large record, 3× on GPT-2, 2.4× on LRA; and the first Transformers above chance on Path-X (16K) and Path-256 (64K)

What came before it

This note named no antecedent, and there are two.

Rabe and Staats (2021), LIT-655 — six months earlier, the same exactness and the same avoidance of the N × N matrix, with O(1) memory per query and a practical O(√n) implementation giving 59× less memory overhead at length 16,384. This paper’s v3 acknowledges them the other way round too: they credit Dao et al. with showing “that the reduced memory requirements can translate to significant speedups on GPUs.”

Jang et al. (2019), unheld here — “lazy softmax”, the identity that lets the denominator be deferred. Rabe and Staats say in as many words that their own Equation 1 is “a rediscovery” of it, while noting that paper discusses neither memory complexity, numerical stability, nor backpropagation.

What this paper contributes is unchanged, and it is the systems half. IO-awareness as the framing, Theorem 2’s HBM-access bound, Proposition 3’s IO-optimality, derived block sizes, and a CUDA kernel. The 3× on GPT-2 is a GPU number, and the antecedent measures the same algorithm on TPU at within a few percent of standard attention — which is the condition SOTA-085 now carries.

Appendix B.5 sets the two side by side: FlashAttention against Rabe and Staats (LIT-655). Both tile the attention matrix and recompute it in the backward pass. Rabe and Staats minimise total memory footprint, while this paper minimises memory accesses, so FlashAttention runs 2–4× faster than standard attention where theirs is “around the same speed or slightly slower”. It also needs less memory: it updates one output incrementally instead of keeping one per block, and derives the backward pass analytically rather than by gradient checkpointing.

Standing in the anthology

The source of SOTA-085, SOTA-086, SOTA-087 and SOTA-114.

Read in full for #114; the reading is NOTE-005. All four practices are confirmed — the second such cluster, after LIT-083.

The paper’s own framing is worth carrying: it argues that a whole literature of approximate attention was optimising a cost model rather than a cost. That is a claim about method, and it generalises past attention.