Most improvements to attention have been approximations: restrict which positions may see which, or replace the mechanism with something cheaper that behaves roughly the same. This one is not. It computes the identical result to ordinary attention , to the last bit, and it is several times faster and uses a fraction of the memory. That combination should be suspicious — you do not usually get exactness and speed together — and the resolution is that the original implementation was not slow because of the arithmetic it performed. It was slow because of the data it moved. This piece is about the difference between those two things, and about the reordering of the computation that makes the largest intermediate result never need to exist.
The bottleneck was never the arithmetic
An accelerator has two kinds of memory. There is a large pool, tens of gigabytes, with high bandwidth by any ordinary standard and yet far too slow to keep the arithmetic units fed. And there is a very small pool — tens of kilobytes per processing unit — sitting immediately next to those units, faster by roughly an order of magnitude.
The ratio between how fast a chip can compute and how fast it can fetch has been widening for decades. A current accelerator performs several hundred arithmetic operations in the time it takes to bring a single byte from the large pool. Consequently the useful measure of an operation is not how many operations it needs but its arithmetic intensity: operations performed per byte moved. Above a few hundred, the chip is limited by its arithmetic and running near its rated speed. Below, it is limited by memory traffic and idling.
Matrix multiplication has high intensity, because each element that gets loaded participates in many products. Almost everything else — adding two tensors, applying a nonlinearity, scaling by a constant, taking a softmax — has intensity close to one. Such an operation reads a value, does a trivial amount of work, and writes it back, spending nearly all its time waiting.
Now look at the standard attention implementation with that lens. It performs a matrix multiply to get the scores, writes the full score matrix to the large pool, reads it back to apply the mask, writes it back, reads it again for the softmax, writes it back, reads it once more to multiply by the values. The two matrix multiplications are efficient; everything between them is memory-bound; and each of those intermediate steps is moving a tensor of size $n^2$ per head.
That tensor is the problem. For a sequence of eight thousand tokens with thirty-two heads, the score matrix alone is over four gigabytes in half precision, per layer. Writing it out and reading it back several times is where the time goes.
Figure 1. Every round trip the standard implementation pays.
Do not materialize the matrix
The idea is to fuse the whole sequence — scores, mask, softmax, weighted sum — into a single pass that never writes the score matrix anywhere, keeping every intermediate in the fast memory beside the arithmetic units.
The obstacle is that the fast memory is small. A score matrix of eight thousand by eight thousand does not fit in tens of kilobytes, not remotely. So the computation is tiled: load a block of queries and a block of keys and values, compute the small block of scores those produce, do everything that needs doing with it, accumulate the contribution to the output, discard the block, and move on. The full matrix is computed, piece by piece, and no piece outlives its usefulness.
The result is a large reduction in traffic to the slow pool — in practice around ninefold, and provably close to the best any exact method can do for a given size of fast memory. It is worth being precise about what falls and what does not: the memory needed becomes linear in sequence length, because the score matrix is never stored, while the traffic stays quadratic in $n$ and is reduced by a large constant factor rather than to a lower order. The arithmetic is unchanged, and indeed the method does slightly more of it than before. It is faster anyway, because arithmetic was never what it was waiting for.
Figure 2. Tiling the score matrix that is never stored.
A pinch of math: making the softmax work in pieces
There is a real obstacle to tiling, and overcoming it is the technical heart of the method.
A softmax is not a local operation. To normalize a row you need its denominator, and the denominator is a sum over the entire row — which means over every key, which is exactly what tiling refuses to have available at once. Processing one block of keys, you cannot yet know what to divide by.
The resolution is to compute the softmax incrementally, maintaining a running result and correcting it as each new block arrives. Keep three quantities per query: the largest score seen so far $m$, the running sum of exponentials $\ell$, and the running weighted sum of values $o$.
When a new block arrives with its own maximum $\tilde{m}$, the running maximum updates to $m^{\text{new}} = \max(m, \tilde{m})$. Everything accumulated so far was exponentiated relative to the old maximum, so it must be rescaled by the factor $e^{\,m - m^{\text{new}}}$:
$$\ell^{\text{new}} = e^{\,m - m^{\text{new}}}\,\ell \;+\; \sum_{j \in \text{block}} e^{\,s_j - m^{\text{new}}}$$$$o^{\text{new}} = e^{\,m - m^{\text{new}}}\,o \;+\; \sum_{j \in \text{block}} e^{\,s_j - m^{\text{new}}}\, v_j$$At the end, divide $o$ by $\ell$ once, and the answer is exactly what the full softmax would have given.
Two things are worth noticing. The rescaling factor is always $e$ to a non-positive power, so it is at most one and the arithmetic stays in a safe range — this is the same maximum-subtraction trick that keeps an ordinary softmax stable, described in the article on logits , applied online. And the correction is exact, not approximate. Nothing is being estimated or truncated. That is why the output matches the standard implementation bit for bit, and why this is a change to the schedule of a computation rather than to the computation.
The backward pass, and where the memory saving comes from
Training needs the gradient, and the gradient with respect to the values requires the attention weights — the very matrix that was never stored. The obvious response is to store it after all, which would surrender the whole gain.
Instead it is recomputed. During the forward pass, keep only the normalization statistics for each query, which is two numbers per query rather than $n$ — negligible. During the backward pass, recompute each block of scores from the queries and keys that are still available, use the saved statistics to reconstruct the correct softmax values immediately without any second pass, and proceed.
This is gradient checkpointing in miniature, and applied here it is a particularly good bargain: the recomputation is a matrix multiply, which is arithmetic-bound and therefore nearly free on hardware that was idling, while the storage it avoids was memory-bound and expensive. The method does more arithmetic and less waiting, which is the right direction on every chip built in the last decade.
The memory consequence is the one that changed what was buildable. Attention’s memory requirement falls from quadratic in sequence length to linear, which removes the constraint that had made long contexts infeasible regardless of how much time you were willing to spend. It is worth being precise about what this does and does not fix: the time is still quadratic, since every pair of positions is still compared. What disappears is the memory wall. That distinction is why the technique enabled long contexts without making them cheap, and why the KV cache rather than the attention matrix is now what limits how many long conversations can be served at once.
Figure 3. What falls, what does not.
The consequence nobody predicted
The most interesting effect of this work was on the research agenda around it.
Before it, a substantial literature was devoted to approximating attention — sparse patterns, low-rank factorizations, kernel methods, hashing schemes — all motivated by the quadratic cost, and all trading some accuracy for tractability. Several were genuinely clever. A good many are no longer used.
The reason is that they had been competing against a badly implemented baseline. Once exact attention became several times faster and linear in memory, an approximation had to beat that to be worth its accuracy cost, and many could not — partly because the approximations were themselves memory-bound operations with poor arithmetic intensity, and partly because on the sequence lengths people actually use, the quadratic term is not yet dominant. The pressure to approximate was, to a large extent, pressure created by an implementation that moved more data than it needed to.
The general lesson is worth extracting, since it recurs. A cost model that counts operations and ignores data movement will systematically misidentify what is expensive, and a research programme built on such a model will optimize the wrong thing. The work that followed — hardware-aware kernels for other operations, fused implementations throughout the stack, the whole practice of writing custom kernels for model internals — descends directly from noticing that once.
What to carry forward
FlashAttention computes exactly the same attention as the standard implementation and is several times faster, which is possible because the original was not limited by arithmetic but by data movement: the score matrix, quadratic in sequence length, was written to slow memory and read back for masking, for the softmax, and for the final multiply, and each of those intermediate steps had an arithmetic intensity near one. The method fuses the whole sequence into a single pass that keeps everything in the small fast memory beside the arithmetic units, tiling over blocks of queries and keys so that the full matrix is computed piece by piece and no piece is ever stored. The obstacle is that a softmax needs a sum over an entire row, which tiling withholds — resolved by an online softmax that carries a running maximum, a running denominator, and a running output, rescaling the accumulated totals whenever a larger score appears, which is exact rather than approximate. The backward pass recomputes the score blocks rather than storing them, saving only two statistics per query, which is a good trade because recomputation is arithmetic-bound and storage was memory-bound. Attention’s memory cost drops from quadratic to linear while both its time and its slow-memory traffic remain quadratic — the traffic reduced by a large constant factor, around ninefold, rather than to a lower order — so what the technique removed was the memory wall rather than the cost — which is why long contexts became possible and remain expensive, and why the cache rather than the attention matrix is now the binding constraint on serving. And its most lasting effect was to obsolete much of the literature on approximating attention, which had been competing against an implementation that moved far more data than it needed to.
The figures in this article follow the shared figure key , which explains what every shape, line style and colour means.
