Who should read it
Core reading for AI-infrastructure, GPU-performance and efficiency-research roles, and valuable for research engineers who train long-context models: it is the clearest example of designing an algorithm around memory traffic rather than FLOPs. Applied-AI and most ML-engineering roles only need the one-paragraph idea (exact attention computed in tiles, never materializing the attention matrix). Skip the block-sparse extension unless you work on sparse attention.
Study objectives
- Explain why standard attention is memory-bound on GPUs and where the N×N matrix is written to and read from high-bandwidth memory.
- Derive the online-softmax rescaling that allows block-by-block computation, and explain recomputation in the backward pass.
- Follow the IO-complexity analysis and the lower-bound statement, and say what they do and do not imply about wall-clock speed.
- Interpret the kernel benchmarks and end-to-end training results, including which baselines, hardware and settings they use.
Prerequisites
Main concepts to understand
- GPU memory hierarchy: large, slow HBM versus small, fast on-chip SRAM
- Compute-bound versus memory-bound operations and kernel fusion
- Tiling Q, K and V into blocks that fit in SRAM
- Online softmax: running maxima and normalizers with rescaling of partial outputs
- Recomputing attention in the backward pass instead of storing the attention matrix
- IO complexity: counting HBM reads and writes as a function of sequence length, head dimension and SRAM size
- Block-sparse FlashAttention as an approximate extension
What to inspect in the evaluation
- Kernel benchmarks (Section 4.3 and the appendix): runtime and memory against standard PyTorch attention and approximate or sparse methods across sequence lengths; check head dimension, batch size, precision, dropout and masking settings.
- End-to-end training (Section 4.1): speedups for BERT and GPT-2 against named implementations; check whether baselines use their best available kernels and identical hyperparameters.
- Quality with longer context (Section 4.2): separate the kernel’s contribution (making long sequences affordable) from the modelling effect of training on longer sequences.
- Hardware dependence: the appendix reports results on different GPUs and configurations, so treat any single speedup as configuration-specific.
- Measurement method: runtimes are averages over repeated calls; note whether variance is reported and whether warm-up is described.
- Exactness: verify the claim with your own comparison against a reference implementation, remembering that floating-point summation order differs.
Guarantees versus empirical results
Correctness or design arguments
- Correctness: the tiled algorithm returns exactly softmax(QKᵀ)V in exact arithmetic (theorem with proof in the appendix).
- Extra memory beyond inputs and outputs grows linearly in sequence length.
- IO analysis: an asymptotic count of HBM accesses for FlashAttention versus standard attention, under the paper’s cost model and SRAM-size assumptions.
- Lower bound: no exact attention algorithm can asymptotically use fewer HBM accesses across all SRAM sizes in the stated range (stated as a proposition).
Empirical results (depend on setup)
- Wall-clock speedups of the attention kernel and of end-to-end training on specific GPUs and model configurations.
- Memory savings measured across sequence lengths.
- Quality gains attributed to the longer context lengths the kernel makes affordable.
- Comparisons with approximate attention methods, which depend on those methods’ implementations and settings.
Reading plan
- Pass 1 (20 minutes): abstract, Figure 1, Section 2.1 on the memory hierarchy and the limitations section; write the main idea in two sentences.
- Pass 2 (90 minutes): Sections 2.2 and 3.1 with Algorithm 1 side by side; derive the rescaling on paper; read the theorem statements in Section 3.2; skim Section 4.
- Pass 3 (3–6 hours, infrastructure and efficiency roles): implement a blocked forward pass in PyTorch that matches a reference, read the backward-pass appendix, and compare with the work-partitioning changes in FlashAttention-2.
Discussion prompts
- Why can an algorithm that performs more FLOPs (recomputation) run faster, and under which hardware conditions would it not?
- What exactly does the IO lower bound say, and why does it not imply wall-clock optimality?
- How does the online softmax relate to the max-subtraction trick used for numerical stability?
- How would you test that a FlashAttention-style kernel is correct in bf16, and which tolerance would you accept?
- Which parts of the speedup carry over to autoregressive decoding with a KV cache, and which do not?