Attention is the operation at the heart of every large language model, and for years the way everyone computed it was quietly wasting almost all of the GPU. The fix, flash attention, made models train and run several times faster, and it did it without changing the math at all, without approximating anything.
↳ It produces the exact same answer as before. The entire speedup comes from one idea: never building a giant thing you were about to throw away. Understanding how it works teaches you the single most important lesson in GPU performance: that moving data, not doing math, is what costs you.
Start with what attention actually needs to do. Every word in a sequence has to look at every other word to decide how much to pay attention to it. So if you have a sequence of some length, you compute a score between every pair of words.
That gives you a big square grid of scores, every word against every word. Then you soften those scores into weights that sum to one, and use them to produce each word's output. Conceptually simple: a grid of pairwise scores, softened, then applied.
Now the trap in how this was done naively. That square grid of scores is enormous. For a long sequence, the number of pairs grows with the square of the length, so doubling the sequence quadruples the grid.
And the old approach built that entire grid, wrote it all out to the GPU's main memory, then read it all back to soften it, wrote it out again, then read it back once more to apply it. The grid was far too big to keep in the small fast memory near the cores, so it lived in the big slow memory, and the algorithm shuttled the whole massive thing back and forth several times.
And here is why that is a catastrophe, tying back to how a GPU really works. The cores can do the math almost instantly. The slow part is fetching from that main memory.
So an algorithm that writes and reads a giant grid several times spends nearly all its time waiting on memory traffic, with the cores idle. The actual computation, the scoring and softening, is trivial. The cost is entirely in hauling that enormous grid in and out of slow memory over and over. The operation was completely memory-bound, bottlenecked on data movement, and the expensive math units mostly sat watching.
So the natural question. If the giant grid is the problem, can we just not store it? Compute it in small pieces in the fast memory near the cores, use each piece immediately, and never write the whole thing to slow memory at all?
That is exactly the idea, and the obstacle is the softening step. Softening scores into weights that sum to one normally needs the whole row of scores at once, because to make them sum to one you have to divide by their total, and you cannot know the total until you have seen every score in the row. So it seems like you are forced to compute and hold the entire row, the entire grid, before you can soften anything. That requirement is what chained everyone to building the full grid.
Flash attention breaks that chain with a running-total trick. Instead of waiting to see the whole row, it processes the row in small blocks that fit in the fast memory. It keeps a running summary as it goes, a partial result and a little bookkeeping about the totals so far.
When the next block arrives, it gently corrects the running result to account for the new information, rescaling what it had so far so that everything stays consistent, as if it had seen those new scores all along. Block by block it folds in each piece, continuously adjusting, and when the last block is done, the running result is exactly, not approximately, what you would have gotten from the whole row at once. The softening happened incrementally, so the full row never needed to exist in one place.
With that unlocked, the whole operation changes character. You load a small block of words into the fast memory, compute their scores right there, fold them into the running result with the trick, and discard them. Then the next block, and the next.
The giant grid is never written to slow memory, never even fully assembled. It exists only as tiny transient pieces in fast memory, each used and thrown away. The several round trips of a huge grid through slow memory collapse into essentially one pass over the actual inputs. The memory traffic plummets, and since memory traffic was the bottleneck, the whole thing speeds up several times over.
Now the objection that reveals the deeper principle. By not storing the grid, you sometimes have to recompute pieces of it later, for instance during the training backward pass, where the old method would have just reread the stored grid. So flash attention deliberately does more math, redoing some scoring rather than saving it. How is doing extra work faster?
Because on a modern GPU, math is nearly free and memory movement is expensive. Recomputing a block of scores in fast memory costs a handful of cheap operations. Rereading that block from slow memory would cost a long, expensive trip that stalls the cores. So trading a big slow memory read for some extra fast math is a winning trade, even though it sounds backwards. This is the counterintuitive truth at the center of modern GPU performance: it is often faster to recompute a value than to store it and fetch it back later. Cheap math beats expensive memory. Flash attention is that principle taken to its logical extreme.
So the whole picture. Attention naively builds a giant grid of every word against every word, shuttling it through slow memory several times, which leaves the fast cores starving and makes the whole operation memory-bound.
Flash attention refuses to build that grid, instead streaming through small blocks in fast memory and folding them into a running result with a trick that lets softening happen incrementally, so the full grid never exists. It even chooses to recompute pieces rather than store them, because doing extra cheap math is faster than paying for slow memory.
The result is the identical answer, several times faster, purely from moving less data. And that is the lesson that outlives this one algorithm...
On a GPU the way to go fast is almost never to compute less.
It is to touch slow memory less...
...even if you have to compute more to get there.