FlashAttention: Attention Was Never Compute-Bound

Imagine a chef with a tiny worktop and a huge pantry down the hall. Every ingredient has to be fetched, and the worktop only holds a few things at once.

If this chef is slow, you might assume they chop too slowly and buy them a better knife. But watch for an hour and you notice they spend almost all their time walking to the pantry and back. The knife was never the problem. The walking was.

For about five years, the AI field bought better knives.

The thing everyone was optimising

Attention is the mechanism that lets a language model relate every word in a sentence to every other word. That "every to every" is where its cost comes from: double the length of the text and the work goes up four times. Everyone knew this and everyone attacked it.

Papers proposed skipping most of the comparisons, or approximating them, or replacing the maths with something cheaper. Dozens reported huge reductions in arithmetic. Almost none of them made anything actually faster on a real chip.

FlashAttention explained why. The bottleneck was never the arithmetic. It was the walking.

See the research paper

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Ré published FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness in May 2022, and it appeared at NeurIPS that year. The word doing the most work in that title is exact.

1. What the chip is actually doing

Attention is three steps. Given the text turned into three sets of numbers — queries, keys and values — for a passage of $N$ words:

$$S = QK^\top \qquad P = \mathrm{softmax}(S) \qquad O = PV$$

In words: compare everything to everything, turn those comparisons into percentages that add to 100, then use those percentages to blend the values together.

The arithmetic in steps one and three is what everyone counted. But watch what the chip has to do to run those three lines.

It computes $S$ — a grid with one entry for every pair of words — and writes the whole grid to memory. It reads the grid back, does the percentages, and writes that grid to memory too. Then it reads it back again and produces the answer.

For a passage of 4,096 words, each of those grids has nearly 17 million entries, for every attention head, in every layer. Written and read in full, twice over.

That is the walking.

2. The pantry and the worktop

A graphics chip does not have one memory. It has at least two, and they differ by a lot in exactly the direction that matters. The paper's own figures, for one common chip:

The worktop is roughly thirteen times faster and about two thousand times smaller. Anything that does not fit on the worktop lives in the pantry, and every trip is paid at the slow rate.

This gives us the idea the paper is built on. For any operation, ask: how much arithmetic does it do per byte it moves?

Do a lot of arithmetic per byte, and the chip's calculators are the limit — cutting arithmetic helps. Do very little, and the calculators sit idle waiting for data — and cutting arithmetic helps not at all.

Now: which category is attention in?

The percentages step is almost pure data movement. So is the masking. So is the scaling. In a standard implementation these cheap steps sit between the two big multiplications, forcing both giant grids out to the pantry and back — and they end up dominating the runtime. The paper measured it: the multiplications account for most of the arithmetic and a minority of the time.

3. The fix: never write the grid down

If the problem is writing giant grids to slow memory, the fix is to not write them at all.

FlashAttention chops the inputs into blocks small enough that a piece of each fits on the worktop together, then loops:

1. Fetch a block of keys and a block of values onto the worktop.
2. For each block of queries: fetch it, compute that patch of the comparison grid on the worktop, do the percentages, blend in the values, and add the result to a running total.
3. Throw the patch away. Never send it to the pantry.

Every entry of the full grid gets computed — all of it. But no entry ever exists in the pantry. It is created on the worktop, used immediately, and discarded.

The pantry traffic drops from something proportional to $N^2$ down to $O(N^2 d^2 / M)$, where $M$ is the worktop size and $d$ is small. That is a large reduction in the thing that was actually costing time.

4. The trick that makes it possible

There is an obstacle, and it is why nobody had done this already.

Turning comparisons into percentages is not a local operation. To work out what share each number deserves, you need the total — and to compute that total safely, you need the largest value too. Both require having seen the entire row. But chopping into blocks means you see each row a piece at a time.

The escape is to keep a running tally and correct it as you go. For each row, hold a running largest value $m$ and a running total $\ell$. When a new block arrives with its own largest value $m'$, rescale what you already have and fold the new piece in:

$$m_{\text{new}} = \max(m, m') \qquad \ell_{\text{new}} = e^{\,m - m_{\text{new}}}\,\ell + e^{\,m' - m_{\text{new}}}\,\ell'$$

The running answer gets rescaled by the same factor. Every correction is exact, so after the last block you have precisely the number you would have got in one pass.

This is the hinge of the whole paper. Without it, chopping attention into blocks would need either two passes or an approximation. With it, one pass is enough and the answer is exactly right.

5. Going backwards: recompute instead of remember

Training needs the comparison grid a second time, to work out how to adjust the model. The forward pass deliberately threw it away.

The obvious answer is to save it — which puts back exactly the memory the design just removed.

FlashAttention instead saves the final answer and a short list of the running tallies — one number per row, not a whole grid — and recomputes the patches it needs on the worktop, block by block.

This does more arithmetic than the standard approach and is still faster. That is the paper's whole thesis restated as an experiment: when the chip is waiting on data, buying extra arithmetic to avoid a trip to the pantry is a good deal.

It also drops attention's memory use from growing with the square of the text length to growing merely in proportion to it — which is what actually made long documents affordable.

6. What it bought, in their numbers

From the abstract:

• 15% faster training of BERT-large at 512 words
• 3× faster on GPT-2 at 1,000 words
• 2.4× on a long-text benchmark at 1,000–4,000 words
• 0.7 better perplexity on GPT-2, from the longer context the memory savings allow
• 6.4 points better on long-document classification

And two results nobody had achieved with this kind of model before: 61.4% on a puzzle called Path-X at 16,000 items, and 63.1% on Path-256 at 64,000.

Those last two are the ones I find most instructive. Before this, these models scored no better than guessing on those tasks — not because they couldn't do them, but because nobody could fit the input in memory to try.

An engineering optimisation did not just make an existing ability cheaper. It produced an ability that did not exist.

7. What came next

The first version removed the memory bottleneck and immediately hit a different one: it was not using the chip's own parallelism well. FlashAttention-2, in 2023, fixed three things — and none of them is about memory traffic at all.

Less non-multiplication work. Modern chips have dedicated hardware for matrix multiplication that is roughly ten times faster than their general-purpose arithmetic. The rescaling in the running-tally trick uses the slow kind. Deferring most of it to the very end removes it from the loop.

Spreading the work across the text. The original split work across separate passages and attention heads. Long-document training has few passages and very long ones, so parts of the chip sat idle. Version 2 splits the text itself.

Better teamwork inside each block, reducing how much the chip's workers have to wait for one another.

FlashAttention-3 goes further for the newest hardware, overlapping the multiplication and percentage stages so the fast units are never idle.

The pattern across all three is worth naming: each version measured, found whatever cost had become dominant, and rearranged around it. The algorithm never changed. The accounting did.

8. The habit worth stealing

The transferable idea is not the block schedule. It is the discipline of measuring what the machine is actually waiting for before optimising anything.

The same reasoning appears everywhere once you look. Systems that treat a model's conversation memory like a filing system do it because serving runs out of space, not arithmetic. Shrinking numbers to fewer bits is usually a data-movement win wearing an arithmetic costume — smaller numbers travel faster, and the unpacking is free because the chip was waiting anyway.

In each case the winning move was to ask which resource is genuinely full, and the answer was rarely the one in the textbook analysis.

Before you optimise, measure how much arithmetic happens per byte moved. If it is low, every calculation you remove is one the machine was not waiting for.

Sources

• Dao, Fu, Ermon, Rudra, Ré — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022
• Milakov, Gimelshein — Online normalizer calculation for softmax, 2018
• Dao — FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023