Flash attention & the memory wall
The bottleneck is memory traffic, not math. Tile the grid through SRAM — same output, 2-4× faster.
Before this
Next in Frontier AI
Stretching the context windowFrontier AI, stop 2 of 102 minThis opens up
Nothing yet. No lesson builds on this one so far.
Transcript
It follows the voice. Click a sentence to jump there.
Attention hands you a bill. Relating every token to every other token means scoring every possible pair — n squared scores for n tokens. You'd assume all that multiplying is the bottleneck. It isn't.
The math units on a GPU are almost never the limit. The chip has a sliver of fast memory called SRAM next to the compute cores — but it holds almost nothing — and a huge pool of HBM — short for high-bandwidth memory — that holds everything, slowly. The wall is the traffic between the two.
By that measure, naive attention is a disaster. It computes all n squared scores and writes them to slow memory, reads them all back for the softmax, then makes another full round trip for the values. Most of the runtime goes to hauling numbers, not computing them.
FlashAttention cuts the grid into tiles small enough to live entirely in SRAM. Each tile gets scored and softmaxed on the spot, with a running max and sum to keep the pieces consistent, and only the finished output leaves.
Tile in, compute, tile out. The full n squared matrix never exists anywhere. Slow memory only handles small tiles going in and finished results coming back.
The output is bit-for-bit identical to the naive version — this is not an approximation. Yet it runs two to four times faster, and the extra memory drops from n squared to n.
So the bill was never the multiplying — it was the hauling. FlashAttention pays it down without changing a single equation. It just moves bytes more carefully. Much of modern AI's speed comes from exactly that: not smarter math, smarter plumbing.