FlashAttention & IO-awareness

01 · Concept

Concept

Qwen3.8-27B accepts a native context of 262,144 tokens, and 16 of its 64 layers run full softmax attention over that entire span. Feed the model a maximal prompt and each of those layers must, in principle, compare every query with every key. The attention formula is one line; the question is whether the intermediate objects it implies can physically exist.

Run the arithmetic. With $N = 262{,}144 = 2^{18}$, the score matrix $QK^\top$ has $2^{18} \times 2^{18} = 2^{36} \approx 6.9 \times 10^{10}$ entries per head. At two bytes each in bf16, that is $2^{37}$ bytes, or 128 GiB, for one head of one layer. Qwen3.8-27B has 24 query heads per full-attention layer, so materializing one layer’s scores would take about 3 TiB, and its 16 full-attention layers together roughly 48 TiB — written to high-bandwidth memory (HBM), read back for softmax, written again as probabilities, and read once more to multiply $V$. An 80 GB-class accelerator cannot hold even a single head’s matrix. Naive attention at this context length is not slow; it is impossible without changing the memory plan.

Accelerators have a memory hierarchy: a small pool of on-chip SRAM measured in tens of megabytes that is very fast, and HBM that is orders of magnitude larger but far slower to reach. Arithmetic units idle while tensors travel between the two. IO-awareness means designing the algorithm around bytes moved through this hierarchy rather than only counting floating-point operations. FlashAttention’s observation is that attention can be computed in tiles that fit SRAM, so the score matrix never needs to exist as a whole.

Split queries into row blocks and keys and values into column blocks, and size the tiles to the head dimension. Qwen’s full-attention heads use head dimension 256, so a 128-row key tile occupies $128 \times 256 \times 2 = 65{,}536$ bytes, 64 KiB, and its value twin another 64 KiB. That is twice the footprint of the 128-dimensional heads many early kernels were tuned for, which forces smaller row blocks or more passes within the same SRAM budget — and it means kernel support for $d = 256$ must be verified for the exact library version, not assumed from an API flag.

The obstacle to tiling is softmax: normalizing a row of scores appears to require seeing the whole row at once. The classic wrong turn is to softmax each tile locally and concatenate the results. Try it on one query row whose scores arrive as a block $[1.0,\ 3.0]$ followed by a block $[5.0]$. A local softmax of the first block gives the score 3.0 a weight of $e^{3}/(e^{1}+e^{3}) \approx 0.881$ — but the true weight, once the dominant 5.0 arrives, is only about 0.117. Local normalization bakes in a denominator that later tiles invalidate.

Online softmax corrects this by carrying two running statistics per row: the maximum $m$ and the exponential sum $\ell$. After the first block, $m_1 = 3$ and $\ell_1 = e^{1-3} + e^{3-3} = 0.135 + 1 = 1.135$. The second block raises the maximum to $m_{new} = 5$, so the old sum is rescaled by $e^{m_1 - m_{new}} = e^{-2}$: $\ell \leftarrow 1.135 \times 0.135 = 0.154$, then the new term adds $e^{5-5} = 1$, giving $\ell = 1.154$. Computing directly, $e^{-4} + e^{-2} + e^{0} = 0.018 + 0.135 + 1 = 1.154$ — identical. Partial output accumulators are rescaled by the same factor, so after the last tile the division by $\ell$ yields exact attention. The operation order changes, so the last floating-point bits may differ, but nothing is approximated, sparsified, or truncated.

Recomputation extends the same logic to training: the backward pass can recompute local scores inside tiles instead of reading a stored $N^2$ matrix, trading cheap arithmetic for expensive HBM traffic.

The benefit is phase-specific. Prefill processes many query rows and is exactly the regime the tiling targets. Single-token decode has one query row; as lesson 7.3 established, it is dominated by reading weights and the KV cache (lesson 7.2 derives Qwen’s per-token cache cost), and specialized decode kernels are cousins rather than beneficiaries of the same headline speedup.

The durable lesson outlives one kernel. When arithmetic is cheap and memory movement is not, algorithm design must count bytes through the hierarchy. FlashAttention is exact attention rearranged around that physical reality — and it is the reason a quarter-million-token context is a product feature rather than a thought experiment.