Exact attention, sped up by never writing the giant grid down. · Source

The same attention math — just never written down.

Attention's real cost isn't the multiplying — it's hauling a giant table of scores in and out of slow memory. Flash attention changes nothing about the math and everything about where it happens. Same answer, a fraction of the time. Like speed-solving a cube: same puzzle, same moves — the win is pure method.

The bottleneck isn't the math. It's the moving.

Standard attention builds the full table of every word scored against every other — then writes that giant grid to slow memory, reads it back to normalize, and reads it again to weight the values. Three round-trips with the biggest object in the model, while the chip finishes the math and just waits. Like an unrolled carpet: too big for the room, so it lives down the hall — and every glance means another trip.

A fast chip, starved by a slow pipe.

T \approx \max\!\left(\dfrac{\text{FLOPs}}{\text{compute rate}},\ \dfrac{\text{bytes moved}}{\text{bandwidth}}\right)

A processor does arithmetic far faster than it can fetch the numbers to chew on. So total time is whichever is slower: doing the math, or moving the bytes. Attention is memory-bound — the math is ready, the data isn't — so cutting trips, not cutting math, is what speeds it up. Like a mighty pump on a thin hose: it could move an ocean, but the pipe decides the flow.

So don't load it all. Work one block at a time.

S_{ij} = \dfrac{Q_i K_j^{\top}}{\sqrt{d}}, \qquad Q_i,\, K_j \in \mathbb{R}^{B\times d}

The fix starts simple: never bring the whole grid close. Slice the Keys and Values into blocks small enough to sit on the chip's tiny fast scratchpad. Pull one block in, do its share of the work, drop it, fetch the next. In plain words: each block of B rows fits in fast memory, so you handle just one block-pair at a time. Like carrying firewood: you don't drag the whole woodpile to the hearth — you bring one armful.

But softmax needs the whole row at once.

\mathrm{softmax}(x)_i = \dfrac{e^{\,x_i - m}}{\sum_{j=1}^{n} e^{\,x_j - m}}, \qquad m = \max_{j} x_j

Here's the snag. To turn scores into weights, softmax divides each by the total over every key — and for safety it first subtracts the row's largest score, so the exponentials never blow up. Both the max and the sum need the entire row. But you're holding just one block. How do you normalize against numbers you haven't seen? Like pouring fair shares from one pitcher: you can't fill the first glass until you know how many there are.

The trick: keep a running max, and rebase as you go.

\begin{aligned} m_{\text{new}} &= \max(m_{\text{old}},\, \tilde m) \\ \ell_{\text{new}} &= e^{\,m_{\text{old}} - m_{\text{new}}}\,\ell_{\text{old}} + \sum_{i} e^{\,x_i - m_{\text{new}}} \end{aligned}

You don't need the whole row first — you need a running tally. Track the biggest score seen so far and the running total. When a new block brings a bigger score, put the old total on the new scale with one shrink factor, then fold the block in; the output rescales the same way. The result is identical to doing it all at once. Like a high-jump bar: raise it when a better jumper arrives, and re-measure every earlier mark against the new height.

Same answer. The giant grid never existed.

\text{memory: } O(n^2)\ \longrightarrow\ O(n), \qquad \mathbf{o}_{\text{flash}} = \mathbf{o}_{\text{exact}}

Put it together and the n×n table is never stored — only a running max, a running sum, and the output, all small. Memory falls from quadratic to linear, and the speed-up comes from far fewer trips to slow memory, not from cutting math. And this is no approximation: the result is exact attention. Like folding an origami crane: one sheet, the finished form, no scraps left on the table.

🌱 Maybe it was never the thinking that was slow.

Flash attention adds no new idea — it just respects where the math lives, rearranging the same thoughts for the body that holds them. So how much of any mind's speed is the thinking itself — and how much is just how fast it can reach what it already knows?