Flash Attention: IO-Awareness Explained Without the Kernel Code

Flash attention performs the same number of floating-point operations as the textbook implementation and produces the same result to numerical tolerance. It is faster because it moves far less data between the chip and its memory. That single idea is the whole technique.

The problem is traffic, not arithmetic

Written directly from the definition, attention for a sequence of S tokens with head dimension d does this: compute S × S scores, write them to memory, read them back, apply a softmax over each row, write that back, read it again, multiply by the values. Four full round trips to device memory for a matrix whose size grows with the square of the sequence length.

The arithmetic is O(S²d) and so is unavoidable. The memory traffic is O(S²) and is entirely an artefact of how the computation was written down. Since attention at modest head dimensions has low arithmetic intensity to begin with, that traffic is what the operation actually costs.

The reason it was written down that way is instructive rather than careless. Expressed as a sequence of standard operations — a matrix multiply, a softmax, another matrix multiply — every step is a library call that is individually well optimised, and each one necessarily reads its input from memory and writes its output back because that is the contract of a library call. The inefficiency lives in the seams between the operations, not inside any of them, which is exactly the class of problem that only a fused implementation can address.

Counting the bytes

Standard, per head, S = 8192, fp16 (2 bytes):

  score matrix S*S = 8192^2 = 67.1e6 elements = 134 MB

  written once, read for softmax, written again,
  read again for the value multiply  ->  ~4 x 134 MB
                                      = ~537 MB per head

Multiply by heads and by layers and the traffic for one
forward pass is measured in hundreds of gigabytes.

Flash attention, same S:

  the S*S matrix is never written to device memory at all.
  Traffic is O(S*d) for Q, K, V and the output, plus a
  re-read of K and V once per query tile.

The peak memory saving is just as important as the speed. The standard form must allocate the S × S matrix, so memory grows quadratically with context — which is why long-context models were impractical before this class of kernel, independent of how long they took.

Tiling, and the softmax obstacle

The fix is the standard one for a memory-bound kernel: tile the computation so each tile fits in the fast on-chip memory, and finish all the work on a tile before moving on. Load a block of queries, then stream blocks of keys and values through it, accumulating the output as you go. The scores for a tile live in on-chip memory and are discarded when the tile is done.

One thing blocks this, and it is the interesting part. Softmax is not local. Normalising a row requires the sum of exponentials over the whole row, and subtracting the row maximum first — the standard trick for numerical stability — requires the maximum over the whole row. You appear to need all the scores before you can normalise any of them, which is exactly what tiling refuses to give you.

The online softmax

The resolution is to keep a running maximum and a running sum, and rescale the accumulated result whenever the maximum changes. Processing a new block:

• Compute the block’s scores and its local maximum.
• Let m_new = max(m_old, m_block).
• Rescale the running sum and the accumulated output by exp(m_old − m_new), correcting for the fact that they were normalised against a maximum that is now out of date.
• Add this block’s contribution, normalised against m_new.

Each correction is a scalar multiply per row, negligible against the matrix work. At the end the running sum is the true denominator and the accumulated output is exactly the attention output. This is the same idea as computing a streaming mean without storing the stream, applied to a softmax — and it is the reason the whole approach is possible at all.

The technique was published as FlashAttention by Dao and colleagues in 2022, with subsequent versions improving the work partitioning across the chip rather than changing this core idea. The name for the general principle, and the more useful thing to remember, is IO-awareness: count the trips to memory, not just the operations.

Worth stating plainly, because it is the part that most often confuses people reading about it: the result is exact. This is not an approximation of attention in the way that sparse or low-rank attention variants are, and it does not trade quality for speed. Reordering a sum and rescaling partial results changes the arithmetic order, so outputs can differ in the last bits the way any floating-point reassociation does, but the mathematics is the same function. That is unusual and it is why the technique was adopted essentially universally rather than becoming one option among several.

The backward pass carries the same idea with one extra wrinkle worth knowing if you ever fine-tune rather than only serve. Gradients need the attention weights that were deliberately never stored, so they are recomputed from the saved statistics during the backward pass. Trading a little redundant arithmetic for a large reduction in memory traffic is the correct bargain whenever a kernel is memory-bound, and it is a pattern that recurs well beyond attention.

What it changed downstream

• Long context became a memory-capacity question rather than a memory-scaling one. With the quadratic term removed, the binding constraint moved to the KV cache, which grows linearly. That is a materially different problem and it is why the interesting long-context work since has been about cache size.
• Prefill got much cheaper. Prefill is where the full attention matrix would have been computed, so this is where the saving lands, and it shows in time to first token on long prompts.
• Decode benefits differently. Generating one token attends over a long history but with a single query row, so there is no large score matrix to avoid. The decode-specific variants of these kernels target the parallelisation of the key/value scan instead.
• It became a portability requirement. Because this kernel is hand-written per architecture, “is there an IO-aware attention implementation for this backend” is one of the sharpest questions to ask about any non-incumbent compute platform.