FlashAttention: IO-Aware Attention

Standard attention (and the $T\times T$ issue)

For one head, with $Q,K,V\in\mathbb R^{T\times d_k}$:

The attention weights matrix is $P\in\mathbb R^{T\times T}$. If you store it explicitly, memory scales like $O(T^2)$.

Online softmax (streaming trick)

For a single row of scores $s\in\mathbb R^T$, the softmax uses:

You can compute this in blocks without ever storing all scores at once by keeping a running maximum $m$ and a running normalizer $\ell$ (Milakov & Gimelshein, 2018).

For a block $b$ of scores with:

you can merge block statistics with the running state via:

At the end, the attention output for that row is:

This is the core idea: you can compute $PV$ without ever storing all of $P$.

One answer, two execution schedules

The optional schedule investigation uses a separate four-token, one-head example, not the size controls or the six-score row above. Its $Q,K,V$ each have shape $T\times d$, with $T=4$ tokens and $d=2$ features. The causal mask allows key $j$ for query $i$ exactly when $j\leq i$. The maximum and the normalizing sum include only allowed keys. A row with no allowed key has no distribution, so it is rejected. Allowing all keys changes the question; changing the tile width must leave the answer unchanged for the same mask and arrays.

Both schedules use the scaled scores and weighted average above. The materialized schedule writes all $T^2$ scores to a slow array store, overwrites them with probabilities, then reads those probabilities to form the output. The tiled schedule holds a query tile and its running numerator, maximum and sum in a fast scratchpad while visiting key/value tiles, and writes only the final output back. Here the query tiles form the outer loop; this simplified schedule is not Algorithm 1 of the paper, whose outer loop runs over key/value blocks.

Three costs answer different questions. Arithmetic counts executed scalar additions, subtractions, multiplications and divisions. Traffic counts numeric words read or written between the slow store and the scratchpad, including repeated reads. Peak live storage counts allocated words that exist at the same time: inputs, output, score/weight buffers, running statistics and workspace. Storing less need not mean fewer transfers, fewer operations or less time.

Let $A$ count allowed query/key pairs, $B$ be the query/key tile width, and $R$ count active query-row/key-tile pairs: pairs of one query row and one key tile that contain at least one allowed key. For these two schedules, total traffic in words is

For the materialized schedule, reads are $Td+T^2d+2T^2+Ad$: queries, keys loaded again for every query, the score and probability rows, and the allowed value rows. Writes are $2T^2+Td$: scores, probabilities and output. For the tiled schedule, reading queries and writing the output contribute $2Td$, and each query tile reads all keys and values. This schedule loads key/value tiles even when some or all of their pairs are masked; skipping those loads would be a different schedule with different counts.

All $Q,K,V,O$ arrays sit in slow storage from the start, including the not-yet-filled output. They occupy $4Td$ words; only the materialized schedule adds a $T^2$ buffer, reused in place for scores and then probabilities. The inputs already exist when counting begins, and reserving output or score storage is not counted as writing. The materialized schedule's largest scratchpad use holds one probability row, one value row, one output row and six scalar workspace words: $T+2d+6$. For a tiled step with $r$ query rows and $c$ key rows, the scratchpad holds

Partial tiles use their actual $r,c$; the peak is the largest of these. The six scalar slots are reused for dot-product sums, scalar products, old and new maxima, the rescaling factor and shifted scores. Built-in operations such as exp count as single steps, and their own internal workspace is not counted. In both schedules, constants, indices, instructions, caches, allocator metadata and the step-by-step display are outside the count. Eight bytes per number is a convention chosen for this example, not a measurement of browser memory or of the paper's precision. The scratchpad holds 32 words. Query and key tiles four tokens wide would need 62 words, so the code refuses that setting with a capacity error rather than spill data it does not count.

The mask is treated as control logic: no mask array is stored or transferred. The Python lists of true/false values show the rule for inspection; an implementation that loads a stored mask would need extra cost entries.

Each dot product starts at zero, counts $d$ multiplications and $d$ additions, then one multiplication by the precomputed $1/\sqrt d$. Exponentials and maximum comparisons are counted separately rather than converted into a FLOP count. Mask tests and indexing are not counted. The materialized arithmetic total is $T^2(2d+1)+A(3+2d)$. The tiled total is $T^2(2d+1)+A(2+2d)+(R-T)+R(d+1)+Td$: the extra terms come from rescaling the running state and from dividing the final numerator. On a row's first active tile the rescaling factor is set to zero, so no exponential of $-\infty$ is evaluated. A tile-row with no allowed key leaves the running state unchanged.

The first active tile still performs the $d+1$ multiplications by that zero factor. Every later active tile computes its rescaling factor once, even when the maximum does not change. Sums start at zero, and each allowed weight is added to the rescaled running sum. Maxima start at $-\infty$ and are compared with every allowed score; the tiled schedule also compares the old and new block maxima for each active tile-row. The maximum-comparison counts are therefore $A$ and $A+R$, and the exponential counts $A$ and $A+R-T$. Starting each maximum at the first score instead would give the same answer with slightly fewer operations.

The running numerator is the sum of value vectors weighted by $\exp(s_{ij}-m_i)$ over the allowed keys seen so far, and the running $\ell_i$ is the sum of those same weights. Rescaling both whenever $m_i$ changes keeps them measured against the current maximum; the numerator is divided by $\ell_i$ only at the end. The two schedules agree to within $10^{-12}+10^{-11}|\text{reference}|$, not bit for bit.

For the causal example, $A=10$; at $B=2$, $R=6$. The materialized and tiled schedules use 150 and 168 arithmetic operations, transfer 132 and 48 words, and peak at 48 and 32 words of slow storage. Their scratchpad peaks are 14 and 30 words, so their combined peaks are both 62 words: which memory level you mean matters. None of these counts predicts elapsed time without evidence about a particular implementation and machine.