FlashAttention

FlashAttention (Dao et al.) is exact softmax attention that never writes the n × n score matrix to HBM. It tiles Q, K and V through on-chip SRAM, keeps a small online-softmax state, and writes only the outputs. Speedups of 2–4× on attention, plus large cuts in activation memory, made long-context training and serving practical.

If batching is about scheduler-level intensity, FlashAttention is about IO-aware algorithms inside one kernel. Same roofline mindset, smaller scope. PagedAttention decides where KV lives. FlashAttention decides how a GPU turns Q, K, V into outputs.

Concepts in this section

Tiling through SRAM

FlashAttention loads blocks of Q, K and V that fit in SRAM, computes partial attention with online softmax state, then moves to the next K/V block. Only the final output O, and in training a little softmax metadata, is written back to HBM.

The backward pass recomputes attention tiles from saved Q, K, V instead of reading a stored n × n map. Extra FLOPs, far fewer HBM bytes. On GPUs that trade is often a win.

At n = 8192, d = 128, FP16, materializing S once is about 128 MB per head. FlashAttention keeps tiles that fit in shared memory and streams K/V. Bytes to HBM become dominated by reading Q/K/V and writing O: roughly O(nd) per head instead of O(n²). FLOPs stay about 4 n² d for the two matmuls. Intensity rises with n, which is why the relative win grows with context length.

Causal masks, padding, RoPE and grouped-query attention change tiling details. They do not change the rule: do not warehouse the score matrix in HBM.

Common mistakes

• FlashAttention ≈ approximate attention. Linformer, Performer and sparse patterns trade exactness for cheaper asymptotics. FA keeps O(n²) FLOPs and exact softmax. It only reduces IO and activation memory.
• If I use FA, I do not need paging. FA is the compute path. Paging is KV capacity and layout. Serving needs both.
• SDPA always means FA. PyTorch dispatches among several backends. Log which one ran.

Key takeaways

• Naive attention writes O(n²) scores to HBM and becomes memory-bound.
• Online softmax plus SRAM tiles give exact attention with far less IO.
• Backward uses recomputation to avoid storing full maps.
• Prefill and decode want different backends. FA2/FA3 change mapping, not the math.

Further learning

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTransformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem …arXiv.org FlashAttention-2Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock n…arXiv.org Online normalizer calculation for softmaxThe Softmax function is ubiquitous in machine learning, multiple previous works suggested faster alternatives for it. In this paper we propose a way to compute classical Softmax with fewer memory accesses and hypothes…arXiv.org

Tri Dao presenting FlashAttention at Stanford MedAI: