FlashAttention,
explained.

The years between the transformer and this rung produced a large literature on making attention cheaper, nearly all of it aimed at the quadratic number of operations — sparsify the pattern, approximate the softmax, low-rank the matrix. This rung's change is to point out that operations were never the bottleneck, and to fix the one that was.

The observation: a standard implementation computes the full score matrix, writes it to the accelerator's main memory, reads it back to apply the softmax, writes it again, reads it again to multiply by the values. On modern hardware, arithmetic is enormously faster than that memory traffic, so the kernel spends most of its time waiting on reads and writes of a matrix it only needs transiently. Attention was slow because it was memory-bound, and every method that traded accuracy for fewer operations was optimising the resource that was not scarce.

The fix is to never build the matrix. Process the sequence in tiles small enough to live in the chip's fast on-die memory, and accumulate the result tile by tile — which requires computing a softmax whose denominator is not known until the end, done by carrying a running maximum and running sum and rescaling the accumulated output as each new tile arrives. Stated in units: the output is the same exact attention, up to floating-point rounding, wall-clock time falls substantially, and memory use drops from quadratic in sequence length to linear, because the quadratic object is never stored.

Self-attention has every token compare itself to every other token, producing an N×N scores matrix. For a sequence of length N, that matrix grows with the square of N — at 8,000 tokens it's 64 million numbers. The standard implementation writes that whole matrix out to memory, reads it back to apply softmax, then reads it a third time to multiply by the values.

The paper's diagnosis: the GPU's arithmetic units were mostly idle, waiting on those reads and writes. Attention was memory-bound, not compute-bound. The field had been optimizing the wrong thing.

A GPU has two very different kinds of memory. HBM is the big main memory — tens of gigabytes, but relatively slow. SRAM is tiny on-chip memory — only kilobytes to a few megabytes, but roughly an order of magnitude faster.

Standard attention parks that giant scores matrix in slow HBM and keeps trudging back to it. FlashAttention's entire strategy is to never put it there in the first place.

Instead of computing the full matrix at once, FlashAttention breaks the work into tiles — small blocks of queries and keys that fit inside fast SRAM. It loads a block, computes its piece of attention entirely in SRAM, and moves on, never writing the intermediate scores back to HBM.

The tricky part is softmax, which normally needs to see a whole row at once to normalize it. FlashAttention uses an online softmax: it processes blocks one at a time and keeps a running total and running maximum, rescaling previous partial results as new blocks arrive — so it gets the exact normalized answer without ever holding the full row. All the attention operations are fused into a single pass over fast memory.

The rescaling trick is the whole paper, so let's do it by hand. Softmax is computed "safely" by subtracting the row's max m before exponentiating (so nothing overflows), then dividing by the sum ℓ:

Both m and ℓ look like they need the whole row. They don't. Take a row of scores s = [2, 1, 3, 0] and process it in two blocks — [2, 1] first, then [3, 0]:

Why does one multiply fix everything? Every old term was e^^s−2 but should have been e^^s−3 — and e^^s−3 = e^^s−2·e^²⁻³. The correction factor is the same for every term, so it can be applied to their sum. That's the update rule the kernel runs per block:

The third line is the punchline: the output accumulator O (the weighted sum of values built so far) is corrected by the same factor. Final answer = O / ℓ. Finishing the example: with m = 3 and ℓ = 1.5530, the weights come out [0.2369, 0.0871, 0.6439, 0.0321] — identical to computing softmax on the whole row at once. No approximation anywhere; the N×N matrix simply never existed.

This is what set FlashAttention apart. Earlier attempts to speed up attention approximated it — sparse patterns, low-rank tricks — trading accuracy for speed. FlashAttention changes only the memory schedule. The result is the same attention, exact up to floating-point rounding; you give up nothing.

The kernel takes Q, K and V for one attention head, each N × d and stored in the GPU's high-bandwidth memory (HBM), plus an on-chip SRAM budget of M elements. It returns the same output O = softmax(QK^^T)V a standard implementation would, with the 1/√d scale, masking and dropout fused in, but never writes the N × N score or probability matrix to HBM. The forward pass:

1. Size the tiles from SRAM. Key and value blocks get B_c = ⌈M/4d⌉ rows and query blocks B_r = min(⌈M/4d⌉, d) rows, so a query block, a key block, a value block, an output block and their B_r × B_c score tile all fit on chip at once.
2. Initialize running state in HBM. An output accumulator O = 0, a running row sum ℓ = 0 and a running row maximum m = −∞, one entry of ℓ and m per query row.
3. Outer loop over key and value blocks. Load K_j and V_j into SRAM once each.
4. Inner loop over query blocks. Load Q_i, O_i, ℓ_i and m_i. Compute the tile S_ij = Q_iK_j^^T, its row maxima, P̃_ij = exp(S_ij − row max) and its row sums. Merge them with the running m and ℓ, rescale the old O_i by ℓ_old·exp(m_old − m_new), add P̃_ijV_j rescaled to the new maximum, divide by the new ℓ, and write O_i, ℓ_i and m_i back. Masking and dropout are applied to the tile while it is still in SRAM.
5. Keep only O(N) extra state. After the last block, O is the exact attention output. What is saved for the backward pass is O plus the statistics ℓ and m (and the random-number state for dropout), not the N × N matrix.

The backward pass recomputes instead of reading

Standard backpropagation through attention reads the stored probability matrix P. FlashAttention never stored it, so it rebuilds each tile on chip: reload Q_i, K_j and V_j, recompute S_ij, and get P_ij = exp(S_ij − m_i)/ℓ_i directly from the saved statistics. From there: dV_j += P_ij^^TdO_i, dP_ij = dO_iV_j^^T, dS_ij = P_ij ∘ (dP_ij − D_i), dQ_i += dS_ijK_j and dK_j += dS_ij^^TQ_i. The vector D_i = rowsum(dO_i ∘ O_i) stands in for the reduction the softmax gradient would otherwise need over a full row of P, and it is computed from N × d matrices. Recomputation costs extra arithmetic, but the backward pass makes the same Θ(N²d²/M) HBM accesses as the forward pass, which is why it beats the standard backward despite doing more FLOPs.

Block-sparse FlashAttention

The paper also gives a sparse version: given a fixed mask over blocks, the kernel skips every tile that is zero. Its IO cost falls in proportion to the fraction s of non-zero blocks, Θ(Nd + N²d²s/M). Unlike the dense kernel this one is an approximation, because it drops attention outside the mask, and it is the version that reached the 64K-token Path-256 task. Both versions are written as a single fused CUDA kernel, which is the practical reason the method has to be re-tuned for each GPU generation.

• N — sequence length; d — head dimension; M — the SRAM available to the kernel, in elements.
• N·d — reading Q, K and V and writing O once, which no method can avoid.
• N² — the score and probability matrices the standard kernel writes out and reads back.
• N²·d²/M — every K/V block is loaded once, but each one triggers a pass that streams Q and O through SRAM, and there are N/B_c = Θ(N·d/M) such passes, each moving Θ(N·d) elements.

The quadratic terms differ by a factor of M/d². For head dimensions of 64 to 128 and roughly 100 KB of SRAM, d² is many times smaller than M, so FlashAttention makes many times fewer HBM accesses. At the extremes: if M reached N·d, everything would fit on chip and the cost would fall to a single pass, Θ(N·d). If M shrank toward d, the bound would grow toward N²·d, worse than the standard kernel, so the method relies on SRAM being large compared with d². The paper also proves a matching lower bound: no exact attention algorithm can make asymptotically fewer HBM accesses than Θ(N²d²/M) over the whole range of SRAM sizes from d to N·d.

The main timings are on NVIDIA A100 GPUs. The paper trained BERT-large on the MLPerf 1.1 setup, GPT-2 small and medium on OpenWebText, and Transformers on the Long Range Arena and the Path-X and Path-256 tasks, against PyTorch, HuggingFace, Megatron-LM, Nvidia's MLPerf code and a range of approximate-attention methods.

The GPT-2 runs reached the same perplexity as the baselines, so the speedups (about 3× against HuggingFace and 1.7× against Megatron-LM) came with no change in model quality. The attention-only row is the paper's central argument in numbers: the recomputing kernel did about 13 percent more arithmetic, yet moved roughly a ninth of the bytes and ran about 5.7 times faster, so runtime followed memory traffic, not FLOPs. Varying the block size told the same story: larger blocks meant fewer HBM accesses and lower runtime until arithmetic took over as the limit. Memory use grew linearly with sequence length, up to about 20 times less than the exact-attention baselines.

The negative result is about scale. Dense FlashAttention is still quadratic in arithmetic, and approximate methods such as Linformer overtook it somewhere between 512 and 1,024 tokens; only the block-sparse version was faster than every approximate method at every length. Longer context did pay off in accuracy: GPT-2 with a 4K context trained with FlashAttention was still faster than Megatron-LM at 1K and reached lower perplexity.

The cost is that this is a hardware-specific kernel, not an algorithm you can express in a framework and have run everywhere. Its advantage comes from the ratio between on-die and main memory bandwidth on a particular accelerator, so it has to be rewritten per architecture, and successive versions exist largely to track new hardware. The paper is explicit that the analysis is about the memory hierarchy.

There is nothing to contest in the result — exactness makes it unusually unarguable — so what is worth noting instead is the opportunity cost the field paid for the wrong diagnosis. Several years of approximate-attention research was aimed at a resource that was not binding, and that line largely stopped being competitive the moment exact attention got fast.

How it held up: it is in everything, and it is the single most-transferable lesson on this site for anyone who optimises anything. Measure which resource is saturated before redesigning the mathematics — the first-order win was available without giving up a single bit of accuracy, and it went unnoticed because the community was reasoning about complexity rather than about the memory hierarchy. Long context became practical as a side effect, since the quadratic memory term was what made it impossible.

FlashAttention became a default building block baked into every major training and inference stack — PyTorch's scaled_dot_product_attention, vLLM, TensorRT-LLM, xFormers. If you've used a model with a long context window, you've almost certainly used it. And it grew into a lineage:

Its deeper lesson is one every systems engineer learns eventually: on modern hardware, data movement is the cost, not arithmetic. The fastest code isn't the one that does the least math — it's the one that moves the least memory. FlashAttention is that principle, applied to the single most important operation in modern AI.

Work each one out before you open it — a number, a prediction or a reason, not a definition. Producing the answer is what makes it stick. Your progress is kept on this device.

Five papers about the bill. Each one buys throughput back from a different part of the stack: the weights, the memory hierarchy, the KV cache, the decode loop.

1. 1Knowledge Distillation2014
2. 2FlashAttention2022
3. 3PagedAttention2023
4. 4Speculative Decoding2023
5. 5Medusa2024

All 110 papers, in 19 ladders →