Vol. XVI · Kernels and Efficiency Plate 4 of 8FlashAttention, tiling, SRAM, HBM traffic, IO-aware, score matrix, online softmax, long contextCE

FlashAttention.

The same arithmetic as attention, done in tiles that never leave the chip — so the N × N matrix is never written.

Plain attention computes S = QKᵀ, writes it to HBM, reads it back for the softmax, writes P, reads it back to multiply by V: an N × N matrix per head crossing HBM several times. At 8 192 tokens and 32 heads that is 4 GB of scores per layer in BF16, and the softmax that touches them is purely memory-bound. Dao and colleagues’ FlashAttention (2022) is “IO-aware”: it loads blocks of Q, K and V into the fast on-chip memory of each multiprocessor (plate XIII.2), computes a tile of scores, applies the online softmax (plate XVI.3) to fold the tile into a running output, and moves on — the scores never reach HBM.

The arithmetic is unchanged, and the result is exact. What changes is the traffic: for sequence length N, head size d and on-chip memory M, the paper shows standard attention moves Θ(Nd + N²) elements and FlashAttention Θ(N²d²/M). The workshop compares the two for the chosen length, head size and SRAM, and the time each takes on an H100 by the roofline.

Below: left, HBM traffic per head against sequence length for standard attention and FlashAttention. Right, time per layer at the chosen length by the roofline.

What you will learn
Workshop
Per head, one sequence of N tokens, forward pass: traffic in BF16 elements from Dao et al.’s analysis (standard: 4Nd + 4N²; FlashAttention: 4Nd + 2N²d²/M with M the SRAM in elements). Time per layer for 32 heads at 989 TFLOP/s and 3.35 TB/s by the roofline. Score memory: N² × 2 bytes × 32 heads.

Never write the scores

At 8 192 tokens with 128-dimensional heads, standard attention moves 11.8 times as many bytes as FlashAttention and, by the roofline, takes 4.7 times as long — 5.2 ms per layer against 1.1 — while its 32 score matrices would occupy 4.3 GB. At 131 072 tokens the scores alone would be 1.1 TB per layer, far beyond any accelerator, which is why long-context training (plate XIV.8) presupposes FlashAttention. At 512 tokens the gap is only 3.9×: short sequences barely need it.

The analysis has a twist: FlashAttention’s traffic grows with d²/M, so it helps most when heads are small relative to on-chip memory; with 64-dimensional heads the saving is 39.7× instead of 11.8×. The algorithm also recomputes the scores in the backward pass instead of storing them — the selective recomputation of plate XIV.2, done inside a kernel. Its successors, FlashAttention-2 and -3, change the work partitioning and use the H100’s asynchronous copy engines and FP8 to approach peak; the arithmetic and the result stay the same.

FlashAttention is exact attention, not an approximation: it computes the same softmax(QKᵀ/√d)V up to floating-point rounding. Sparse and linear attention (plate VIII.8) change the computation; FlashAttention changes only where it happens.

Terms

IO-aware
Designed around the cost of moving data between memory levels, not just arithmetic.
Tiling
Splitting a computation into blocks that fit in fast on-chip memory.
SRAM
The fast on-chip memory of each multiprocessor, about 228 KB on an H100.
Score matrix
The N × N matrix QKᵀ that standard attention stores.
FlashAttention keeps attention’s tiles on chip and never writes the N × N scores: the same exact result with far less HBM traffic and memory linear in N.

Sources

  1. Dao, Fu, Ermon, Rudra and Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”, NeurIPS 2022, Theorem 2.
  2. Dao, “FlashAttention-2”, ICLR 2024; Shah et al., “FlashAttention-3”, NeurIPS 2024.