The idea: attention is limited by memory, not by math
Attention computes O = softmax(Q·Kᵀ / √d)·V (see Self-Attention).
The middle step makes an N×N matrix of scores, one for every pair of tokens. A GPU has two
kinds of memory: HBM, tens of gigabytes at about 1.5–2 TB/s, and SRAM on the
chip, only 192 KB per streaming multiprocessor but about 10 times faster. The arithmetic is quick; moving the
N×N matrix in and out of HBM is what takes the time. FlashAttention never writes that matrix:
it works on tiles small enough to stay in SRAM.
The canvas uses one head with N = 8 tokens and d = 4, in fp16 (2 bytes per number), and gives SRAM room for 96 numbers. The top panel is HBM, the bottom left is SRAM, and the bottom right shows the softmax state and the bytes counted so far.
Naive attention: three kernels
A plain implementation runs one GPU kernel per step, and each kernel reads its inputs from HBM and writes its whole result back. Kernel 1 writes S, kernel 2 reads S and writes P, kernel 3 reads P. At real sizes there is no choice: S for N = 4,096 is 32 MiB per head, and SRAM is 192 KB. (Demo: naive attention — S and P round-trip through HBM)
| kernel | reads from HBM | writes to HBM | numbers (N = 8, d = 4) |
|---|---|---|---|
| 1: S = Q·Kᵀ/√d | Q, K: 2Nd | S: N² | 64 + 64 |
| 2: P = softmax(S) | S: N² | P: N² | 64 + 64 |
| 3: O = P·V | P, V: N² + Nd | O: Nd | 96 + 32 |
| total | 2N² + 3Nd | 2N² + Nd | 384 numbers = 768 bytes |
Tiling: one S tile at a time
FlashAttention (in the FlashAttention-2 loop order) reads a tile of B rows of Q into SRAM and keeps it there. The K and V tiles then stream past it one by one. For each pair it computes a B×B tile of S, uses it, and throws it away. When the last K,V tile has passed, the B output rows are finished and written once. (Demo: FlashAttention — tiles never leave SRAM)
Q is read once and O is written once. K and V are read again for every Q tile, N/B times. So the traffic is
2Nd + 2N²d/B numbers instead of 4N² + 4Nd, and no N×N matrix is stored anywhere.
Online softmax: running max and running sum
Softmax divides by a sum over the whole row, and for safety it first subtracts the row's maximum (see Softmax & Cross-Entropy). A tile only sees B scores of the row. The trick is to keep, for every row, a running max m, a running sum ℓ and a running output O. When a new tile arrives:
m_new = max(m_old, max of the tile's scores)
ℓ = ℓ · e^(m_old − m_new) + Σ e^(s − m_new)
O = O · e^(m_old − m_new) + Σ e^(s − m_new) · v
at the end: O = O / ℓIf the new tile has a bigger score, everything summed so far was measured against a max that was too small,
and multiplying by e^(m_old − m_new) (less than 1) corrects it. If the max stays the same the
factor is 1. (Demo: online softmax — a bigger max arrives, rescale)
The answer is exact, not an approximation: after the last tile, m and ℓ equal the max and the sum over the full row. The canvas checks that O matches a full softmax. (Demo: same output, half the bytes)
Counting the bytes
| case | naive | FlashAttention |
|---|---|---|
| canvas: N = 8, d = 4, tiles of 4 | 768 B | 384 B |
| canvas: tiles of 2 | 768 B | 640 B |
| N = 1,024, d = 64, tiles of 128 | 8.5 MiB | 2.3 MiB |
| N = 4,096, d = 64, tiles of 128 | 130 MiB | 33 MiB |
| N = 16,384, d = 64, tiles of 128 | 2 GiB | 516 MiB |
These are per head, per sequence, forward pass only. Real models also apply a mask and dropout to S, and in a naive implementation each of those is one more N×N round trip. The FlashAttention paper measured GPT-2 medium (N = 1,024): 40.3 GB of HBM traffic for standard attention against 4.4 GB, and 41.7 ms against 7.3 ms. The memory win matters as much: the extra memory grows with N, not N², which is what made long contexts practical.
Tile size and SRAM
Bigger tiles mean K and V are read fewer times, but every tile must fit in SRAM together: Q, K, V and O tiles (4·B·d numbers), an S tile (B²) and m, ℓ (2B). On the canvas, tiles of 4 need 88 of the 96 numbers; tiles of 2 need 40 but read K and V 4 times (Demo: tiles of 2); tiles of 8 would need 208 and are refused (Demo: tiles of 8 do not fit in SRAM). Real kernels pick B from the SRAM size, roughly M / 4d, which is why the paper's IO count is about N²d²/M. The same idea, keep the working set in the fast level, runs through the whole memory hierarchy.
The backward pass
Training needs P again to compute gradients. FlashAttention does not store it: it keeps only O and, per row, m and ℓ, and recomputes the S tiles during the backward pass. That costs more arithmetic but much less memory traffic, so it is still faster.
FlashAttention speeds up the attention over the tokens a model has. During text generation the keys and values of earlier tokens are kept in a KV cache, and serving systems store that cache in blocks with PagedAttention; that page shows how. How a Large Language Model Works shows where attention sits in the whole model.
What the page leaves out
Many heads and a batch running in parallel on many SMs, causal masking (FlashAttention skips tiles that are fully masked), dropout, the backward pass in detail, the GPU's L2 cache (which catches some of the K,V re-reads), warps and tensor cores, FlashAttention-3's asynchronous loads and FP8 on Hopper GPUs, and the √d scaling folded into Q. The score matrix on the canvas fits in SRAM; the point is what happens when it does not.