The Attention Matrix You Never Write Down

Before you read

The paper assumes you can read pseudocode for a GPU kernel and that you hold two facts about hardware that machine-learning courses rarely teach. First, a modern accelerator has a memory hierarchy: a large, slow pool of high-bandwidth memory (HBM, the number on the spec sheet) and a tiny, fast scratchpad of on-chip SRAM local to each compute unit. Second, an operation's wall-clock time is often set not by how many arithmetic operations it does but by how many bytes it moves between those two pools. If both are familiar, the paper's contribution will land the first time. If not, the next section builds them from the ground up before the algorithm needs them.

The bottleneck was never the multiply

Self-attention scales quadratically in sequence length: for N tokens and head dimension d, the scores S = QK^⊤ form an N×N matrix, and both the arithmetic and the memory grow as N². For years the field read that quadratic as a compute problem and answered it with approximations: sparse patterns, low-rank factorizations, kernelized attention, each cutting the number of floating-point operations. FlashAttention starts from a different reading of the same fact. On current hardware the quadratic matrix is expensive mostly because it is written to and read from slow memory, not because the multiplications are hard. The fix is therefore not to do less arithmetic but to move fewer bytes, and to do it without changing the answer.¹

The paper names the missing design principle directly, and it is the sentence the rest of the work exists to cash out.¹

Why the score matrix is the tax

Start with the hardware asymmetry the whole idea turns on. On an A100, the paper's reference GPU, arithmetic throughput has outrun memory bandwidth by a wide margin, and most Transformer operations feel that gap.¹

The numbers make the asymmetry concrete. The A100's high-bandwidth memory delivers roughly 1.5 to 2.0 terabytes per second, while the on-chip SRAM local to each of its 108 compute units runs at an estimated 19 terabytes per second, about an order of magnitude faster. That SRAM is also small: 192 kilobytes per unit.¹ The same asymmetry drives every well-tuned GPU kernel; the standard matrix-multiply playbook is to stage small tiles of the operands into SRAM and reuse them there before touching slow memory again.²

Now trace what a textbook attention kernel does with that hierarchy. It computes the full scores S = QK^⊤ and writes all N² of them to HBM. It reads them back to apply the softmax row by row, producing probabilities P, and writes P back. It reads P again to multiply by V. The arithmetic is a pair of matrix multiplies and an elementwise exponential; the expensive part is the round trips of an N-by-N object across the slow bus. When d is a hundred or so and N is in the thousands, that traffic, not the multiplies, sets the clock. This is the same lesson practitioners learn when they profile a model and find it starved for bandwidth rather than compute: fusing operations to avoid materializing intermediates is the single highest-leverage move.³ FlashAttention is that move carried out for the attention block, all the way.¹

The obstacle: softmax refuses to be split

The obvious fix is fusion: never write S or P to HBM at all. Load a block of queries and a block of keys into SRAM, compute their scores there, turn them into the corresponding slice of output, and keep only the small output block. Matrix multiply tiles this way without trouble. Softmax does not, and understanding exactly why is the crux of the paper.

Softmax over a row of scores is not a local operation. To normalize even the first entry you need the sum of exponentials over the entire row, and to compute those exponentials safely you first subtract the row's maximum so that e^x never overflows. Both the maximum and the sum are reductions over all N keys. If you only ever hold one block of keys in SRAM at a time, you cannot know the global maximum or the global denominator when you process the first block. The naive tiled kernel would have to make a full pass to find the maximum, another to accumulate the denominator, and a third to form the output, which puts the N-by-N scores back on the slow bus and defeats the purpose. The quadratic memory traffic the paper wants to eliminate is, at first glance, forced by softmax's global dependency.¹

The running-statistics trick

The escape is a recurrence that computes softmax in a single streaming pass, correcting itself as it goes. It predates FlashAttention: NVIDIA researchers introduced the online-softmax formulation precisely to cut softmax's memory accesses, computing the normalizer while scanning the scores rather than in a separate pass.⁴ FlashAttention's contribution is to fold the value-weighted sum into the same recurrence, so the output itself is produced online.¹

Here is the mechanism, worth re-deriving once by hand. Process the keys in blocks. Carry three running quantities for each query: the maximum score seen so far m, the running denominator ℓ (a sum of exponentials, all measured relative to the current m), and the running output accumulator o. When a new block arrives, compute its local scores and local maximum, then set the new global maximum m' = max(m, local max). The old running sums were exponentiated against the old, smaller maximum, so they are now too large by exactly a factor of e^{(m − m')}. Multiply both the old ℓ and the old o by that single correction factor, then add the new block's contribution, exponentiated against m'. After the last block, divide o by ℓ. The result is bit-for-bit the same softmax attention you would have computed with the whole row in hand.

1. Score the block in SRAM

   With a query block and a key block resident on-chip, compute their scores and the block's local maximum. Nothing touches HBM.¹

2. Update the running maximum

   Set m' = max(m, local max). Every quantity you already accumulated was normalized against the old m, so it now carries a known, correctable error.

3. Rescale, then accumulate

   Multiply the running denominator and the running output by e^{(m − m')}, then add this block's exponentiated scores and its value-weighted sum. The old and new contributions are now on the same footing.

4. Normalize at the end

   After the final block, divide the accumulated output by the accumulated denominator. Exactness is preserved throughout; no approximation enters.¹

The paper states the property that makes this legal in one line, and it is the load-bearing sentence of the method.¹

Tiling forward, recomputing backward

With the recurrence in hand, the forward pass is a doubly nested loop over blocks. The algorithm splits the queries, keys, and values into blocks sized to fit in SRAM, streams them on-chip, and runs the running-statistics update, writing only the final output and a small per-row normalization statistic back to HBM. The N-by-N score matrix is never materialized in slow memory at all.¹

The backward pass poses the mirror-image problem. Gradients through attention need the probabilities P, and a conventional implementation would have stashed that N-by-N matrix during the forward pass to reuse it now. FlashAttention refuses to store it, so it does the arithmetic twice instead. Having kept the output and the softmax normalization statistics, it reconstructs the needed blocks of S and P on-chip during the backward pass, trading extra floating-point work for memory it never has to touch.¹

The recomputation is a form of gradient checkpointing, the technique that trades an extra forward pass for the memory of stored activations.⁵ Here that trade has a sharper justification than usual. Ordinary checkpointing recomputes to save memory and accepts a speed penalty. Here the recomputation is often faster in wall-clock terms, because the redone arithmetic runs against fast SRAM while the alternative would stall on HBM. The same insight that self-attention need not store its quadratic intermediate appeared in closely related work reducing attention's memory to sublinear in sequence length; FlashAttention pairs that memory argument with an IO-aware kernel that also wins on speed.⁶

What it buys, and where the analysis bites

The payoff is stated as a theorem about memory traffic, which is the honest currency here rather than a raw speedup on one benchmark.¹

Fewer HBM accesses translate into the wall-clock results, which is the point the approximate-attention literature kept missing: the paper reports a 7.6× speedup on the attention computation itself on GPT-2, and end-to-end, a 15% training speedup on BERT-large against the MLPerf 1.1 record, 3× on GPT-2, and 2.4× on Long Range Arena,¹ the benchmark suite built to test efficient Transformers on sequences of a thousand to sixteen thousand tokens.⁷ Because memory now grows linearly in sequence length rather than quadratically, longer contexts become reachable, and the block-sparse variant is the first attention to clear better-than-chance on the 16K-token Path-X and 64K-token Path-256 tasks.¹

Read as a reviewer, the strengths and the caveats are both specific. The result is exact, which is the quiet triumph: unlike the sparse and low-rank methods it displaces, FlashAttention changes the runtime profile without changing the model, so there is no accuracy tax to weigh against the speed. The IO-complexity analysis gives the speedup a principled account instead of leaving it to a benchmark. The costs are equally concrete. The whole argument assumes the head dimension d is small enough that a d-wide block of queries, keys, and values fits in a 192-kilobyte scratchpad; at very large head dimensions the tiling loosens and the advantage narrows. The kernel is hand-written CUDA against a specific memory hierarchy, so its benefit is hardware-shaped, not a free property of the math. And the original release under-uses the GPU's arithmetic units: the follow-up, FlashAttention-2, reworks the work partitioning and parallelism to reach 50–73% of the A100's theoretical peak and roughly double the throughput, evidence that the first version left real performance on the table even as it proved the principle.⁸

Read it yourself

A full read adds the two things a reconstruction has to compress: the exact algorithm blocks with their index bookkeeping, and the IO-complexity proof that justifies every design choice.

• read Section 3.1, Algorithm 1: the forward pass in full. Trace the running m, ℓ, and output update through two blocks by hand and the method is yours.
• read Section 3.2, Theorem 2 and its proof: where the d²/M factor comes from, and why FlashAttention is optimal across a range of SRAM sizes.
• skim Section 4, experiments: stare at Figure 1 and the runtime tables. Note which speedups are attention-only and which are end-to-end before quoting either.
• skip Appendix block-sparse derivations: valuable, but the exact dense algorithm is the idea; the sparse extension can wait for a second pass.

One thing first-time readers miss: FlashAttention is not an approximation. The word “exact” in the title is the whole argument. Every rival it benchmarks against changes the outputs; this one changes only the order in which the same outputs are computed.