Why architecture matters here
To understand why FlashAttention matters you have to understand the GPU memory hierarchy, because that is the entire game. A modern GPU has two relevant tiers of memory. HBM (high-bandwidth memory) is the large pool — tens of gigabytes — where your tensors live; it is called high-bandwidth because it is fast relative to CPU DRAM, but it is still far from the compute units and every access to it costs precious time and energy. SRAM (the on-chip shared memory and registers) is tiny — kilobytes per streaming multiprocessor — but roughly an order of magnitude faster and sitting right next to the arithmetic units. The fundamental fact of GPU performance is that most kernels are memory-bound: the arithmetic units sit idle waiting for data to arrive from HBM. Attention is a textbook example.
Standard attention is memory-bound because it round-trips the huge score matrix through HBM multiple times: write the N×N scores, read them back to apply the softmax, write the softmax result, read it back to multiply by V. Each of those passes moves quadratic bytes across the slow HBM boundary. The actual matrix multiplications are fast; the kernel spends most of its time moving the intermediate matrix in and out of HBM. FlashAttention's contribution is to recognize that this intermediate never needs to leave the chip: if you tile the computation so each block of the score matrix is produced, consumed, and thrown away entirely within SRAM, you eliminate the quadratic HBM traffic. The speedup — often two to four times — comes almost entirely from that reduction in memory movement, not from doing less arithmetic.
The second reason it matters is the memory footprint, which is what unlocked long context. Standard attention requires O(N²) memory to hold the score matrix, so a 32k-token context needs a billion-entry matrix per head — physically impossible to hold on a GPU. FlashAttention needs only O(N) memory because it holds one tile plus the running softmax statistics at any moment. This is why the jump from 2k to 128k+ context windows in production models coincided with the adoption of FlashAttention-style kernels: the quadratic memory wall, not the quadratic compute, was the binding constraint, and tiling removed it.
The third reason is that FlashAttention is exact, and this is what separates it from the long line of approximate-attention methods (sparse, low-rank, linearized) that preceded it. Those methods reduced the quadratic cost by changing what attention computes, trading accuracy for speed and forcing model authors to accept a different operation. FlashAttention changes only how the exact same operation is computed, so it drops into any existing transformer with no change to the math or the trained weights. That is why it became the default rather than a niche option: it is a pure systems win with no modeling cost.
The architecture: every piece explained
The top row traces the forward pass. Q, K, and V — the query, key, and value projections — live in HBM. The tile loader streams a block of query rows (say 128 queries) and a block of key/value columns into SRAM. On-chip, the kernel computes QKᵀ for that block — a small block-of-queries × block-of-keys score submatrix — entirely in fast memory. Then the online softmax updates the per-query running statistics: for each query row it tracks the running maximum score seen so far (for numerical stability) and the running sum of exponentials (the softmax denominator).
The middle row is the accumulation. The block's softmax-weighted scores are multiplied by the block's V and added to the query rows' output accumulator, which also lives in SRAM. The subtlety — the rescale step — is that when a later block contains a larger score than any seen before, the running maximum changes, which means every exponential computed against the old maximum was scaled wrong. The online-softmax algorithm corrects this by multiplying the accumulated output and the running sum by a correction factor exp(old_max − new_max), so the final result is exactly as if all scores had been normalized against the true row maximum. This is the mathematical guarantee of exactness. Only after the last key/value block does the kernel write the output tile O back to HBM — one write, not the multiple round-trips of standard attention.
The right of the middle row is the saved statistics. For the backward pass, FlashAttention saves the per-row softmax normalizers (the log-sum-exp, combining the running max m and sum ℓ) — a tiny O(N) amount — rather than the full N×N attention matrix. During backpropagation it recomputes the attention blocks on the fly from Q, K, V and those saved normalizers. This is the classic recomputation-versus-memory trade: spending a little extra compute to avoid storing (and moving) the giant matrix, which is a net win precisely because the kernel is memory-bound.
The bottom row is the memory hierarchy the whole design is organized around. SRAM is kilobytes but effectively terabytes-per-second and on-chip; it holds the current tiles and accumulators. HBM is gigabytes and fast in absolute terms but far from compute; Q, K, V, and the final O live there, and the entire point is to touch it as few times as possible. The ops strip — head dimension, block size, causal masking, numerical stability — is where kernel authors and users tune the tiling to a specific GPU and model shape.
End-to-end flow
Trace one attention layer's forward pass over a sequence of length N with head dimension d. The kernel launches a grid of thread blocks, one per block of query rows. Take a single query-row block of 128 rows. Its output accumulator O_block (128 × d) starts at zero, its running max m starts at negative infinity, and its running sum ℓ starts at zero — all in SRAM.
The kernel loops over the key/value blocks. For the first K/V block it loads those keys and values into SRAM, computes the 128 × block_k score submatrix S = QKᵀ / √d on-chip, finds each query row's max score in this block, updates m, computes exp(S − m), updates ℓ with the block's exponential sum, and accumulates exp(S − m) · V_block into O_block. It then moves to the next K/V block. If that block contains a larger score, m increases; before accumulating, the kernel multiplies the existing O_block and ℓ by exp(old_m − new_m) so everything is now consistently scaled against the new maximum. This correction is what makes the incremental computation exact.
After the last K/V block, each query row's true softmax denominator is exactly ℓ, so the kernel divides O_block by ℓ to get the final attention output, and writes those 128 output rows to HBM once. It also writes the per-row log-sum-exp statistic. The whole score submatrix for this query block was produced and consumed entirely in SRAM; the N×N matrix never existed anywhere. Because the kernel is causal (in a decoder), it also skips K/V blocks that are entirely in the future relative to the query block — a masking optimization that roughly halves the work for causal attention.
For the backward pass, gradients flow back into Q, K, and V. Rather than reading a stored attention matrix — which was never saved — the kernel recomputes each block's attention weights from Q, K, V and the saved per-row log-sum-exp, then computes the block's contribution to the Q, K, and V gradients, again tiling through SRAM and writing only the gradient tensors to HBM. The recomputation adds arithmetic but avoids storing and re-reading the quadratic matrix, and because the operation is memory-bound the trade is favorable. The net effect across forward and backward is a kernel that is both faster in wall-clock time and linear in memory, which is precisely why it displaced the naive implementation everywhere.