FlashAttention and IO-aware exact attention

Why it matters

Attention dominates memory and time for long sequences. FlashAttention computes exact attention with memory linear in sequence length and far less HBM traffic by restructuring the computation so the full attention matrix never leaves on-chip memory. It is the standard example that bytes moved, not FLOPs, often decide speed, and its ideas carry over to the fused attention kernels used for training and prefill today.

Core ideas

Standard attention moves too much data. Per head, the naive implementation computes S = QKᵀ (an s×s matrix), writes it to HBM, reads it back for masking and softmax, writes P, then reads P and V to form O. Memory is O(s²) per head, and the masking and softmax steps are memory-bound passes over that matrix.

Online softmax. Softmax needs a row's maximum and sum, which seems to require the whole row first. Instead, process keys in blocks and keep, for each query row, a running maximum m, a running normalizer ℓ and an unnormalized output o. When a block arrives with maximum m_b, set m' = max(m, m_b), rescale ℓ and o by exp(m − m'), add the block's exp(scores − m') terms, and continue. Dividing o by ℓ at the end gives exactly the softmax-weighted sum, up to floating-point rounding.

Tiling and fusion. Load a block of queries into shared memory (SRAM), stream key and value blocks through it, update the running statistics, and write only the output block and the per-row softmax statistics back to HBM (FlashAttention-2 keeps a single log-sum-exp per row). One fused kernel replaces several passes, and the s×s matrix is never stored.

Backward by recomputation. Instead of keeping P for the backward pass, recompute the needed blocks from Q, K, V and the saved per-row statistics. This adds FLOPs but saves much more memory traffic, so the backward pass also gets faster.

What changes and what does not. FLOPs stay quadratic in sequence length (slightly higher with recomputation). HBM reads and writes drop sharply and attention memory becomes linear in s. The paper's IO-complexity analysis states this precisely in terms of SRAM size. FlashAttention-2 keeps the algorithm and improves how work is split across thread blocks and warps, including parallelism over the sequence dimension.

Common misconceptions

• "FlashAttention approximates attention." It computes exact attention; sparse and low-rank methods are the approximate family.
• "It reduces FLOPs or makes attention subquadratic." FLOPs stay quadratic; the gain comes from avoiding HBM round trips of the s×s matrix.
• "It speeds up decoding as much as training." The largest gains are in training and prefill. Decoding one query per sequence is dominated by reading the KV cache, which calls for different parallelization.
• "The memory savings come from lower precision." They come from not storing the attention matrix; precision is an independent choice.

How to approach it

Write naive attention, then a blockwise version with the online softmax in plain PyTorch, and check it against the reference on random inputs, with and without a causal mask. Then read the paper's algorithm and IO analysis with your implementation open. For depth, study a fused kernel (the Triton tutorial or the CS336 systems assignment) and measure memory and time against the naive version as sequence length grows.