FlashAttention and IO-Aware Attention

Standard attention is usually described purely in terms of arithmetic: multiply queries by keys, take a softmax, multiply by values. But on real hardware, arithmetic is often not the bottleneck — moving data between memory tiers is. A GPU has a small amount of extremely fast on-chip memory (SRAM) and a much larger pool of slower memory (HBM), and standard attention implementations write the full, huge attention matrix out to slow memory and read it back repeatedly. FlashAttention is not a new algorithm in the mathematical sense — it computes the exact same numbers — it is a new way of scheduling the computation so slow-memory traffic collapses, turning attention from a memory-bandwidth-bound operation into something close to compute-bound.

The analogy: cooking with a tiny counter and a distant pantry

Imagine a kitchen where the counter (fast SRAM) fits only a few ingredients at a time, and the pantry (slow HBM) holds everything but is a long walk away. A naive cook doing attention makes a giant $n \times n$ shopping list, walks it to the pantry to write it down, and later walks back to read it again — most of the time is spent walking, not cooking. FlashAttention reorganises the recipe: work on small batches that fit entirely on the counter, finish everything possible with what's there, and only walk to the pantry for the next batch — never writing the full list down at all. The final dish is identical; only the number of trips to the pantry changes.

What standard attention writes to slow memory

$$S = \frac{QK^\top}{\sqrt{d_k}}, \qquad A = \text{softmax}(S), \qquad O = AV$$

In words: standard attention computes a full similarity matrix $S$ of size $n \times n$, turns it into weights $A$ via softmax, then multiplies by values $V$ to get output $O$. $S$ and $A$ are both $n \times n$ — for $n=8{,}192$ tokens, that's over 67 million entries, materialised in slow memory and read back at least twice.

The trick — tiling and online softmax. Split $Q$, $K$, $V$ into small blocks that fit in fast SRAM. Process block pairs one at a time, computing partial output and keeping a running (online) estimate of the softmax normaliser, updated incrementally as each new key/value block arrives — so the full matrix $A$ is never stored at once, only small tiles, transiently, in fast memory.

Worked example 1: counting memory traffic on a small case

Take $n = 4$ tokens, block size $b = 2$. Standard attention: materialise the full $4\times4$ score matrix (16 numbers) in HBM, read it back for softmax, read it again to multiply by $V$ — roughly $3 \times 16 = 48$ "number movements" to/from slow memory for the matrix alone.

FlashAttention with tiling: process key/value block 1 (2 keys) against all 4 queries — a $4\times2$ tile (8 numbers) computed and consumed entirely in SRAM, only a small running max/sum (4 numbers) written to HBM per query. Block 2, same again. Total slow-memory traffic for the intermediate matrix: close to zero — only the small $O(n)$ running statistics ever touch HBM, versus $O(n^2)$ for the naive approach. At $n=4$ the gap looks modest; at $n=8{,}192$ the naive approach moves gigabytes while FlashAttention moves megabytes.

Worked example 2: wall-clock speedup at scale

Reported benchmarks at sequence length 2,048 on a mid-sized transformer block show standard attention taking roughly 41 ms per forward pass, dominated by memory traffic, versus roughly 7 ms for FlashAttention doing the identical computation — about a 5–6× wall-clock speedup, purely from restructuring memory access, output matching to floating-point precision.

Sequence length	Standard attention	FlashAttention
1,024	12 ms	3 ms
2,048	41 ms	7 ms
4,096	156 ms	15 ms

The gap widens as sequence length grows: naive $O(n^2)$ memory traffic outpaces FlashAttention's, which stays close to linear for a fixed block size.

The naive approach's slow-memory traffic grows like the curve above — quadratically — while FlashAttention's traffic grows close to linearly, which is why the gap widens so sharply at long context.

What this means in practice

FlashAttention (and its successors, FlashAttention-2 and -3) is the default attention implementation in essentially every modern training and inference stack, since it makes long-context training and serving feasible on hardware that would otherwise be memory-bandwidth-starved. It composes cleanly with Multi-Query and Grouped-Query Attention — one shrinks the cache, the other shrinks memory traffic during the computation itself.

Practice

1. Why does writing the full $n \times n$ attention matrix to HBM become more costly, relative to the useful compute, as $n$ grows — even though both the compute and the matrix size grow quadratically?
2. Explain in one sentence why FlashAttention and MQA/GQA solve two genuinely different bottlenecks rather than being alternative solutions to the same problem.