The last post ended with vectors. This is the thing those vectors do: attention, the mechanism that made transformers beat everything else. And it's the reason inference has the shape it does.
If I had to keep one fact from this whole series it would be the one below. Almost every serving problem I deal with is a downstream consequence of attention being quadratic, and I did not appreciate how far downstream that reaches until I started tracing problems back to it.
The operation
Attention answers one question: how much should each token look at every other token? For a sequence of N tokens, it computes an N×N matrix of scores, where entry (i, j) says "how much does token i depend on token j".
Concretely, each token produces a query, a key, and a value (three vectors, learned projections). The query of token i is compared to the key of token j to get the score. Then the values are blended by those scores. That's it, that's attention.
Why O(n²) is the whole story
The N×N matrix is the reason inference has the shape it does. It's why:
- Prefill is compute-bound: for long prompts, the N×N matrix is huge, and computing it is real FLOPs.
- The KV cache exists: you need the keys and values of all previous tokens to compute the next token's attention. That's the cache.
- FlashAttention works: the matrix is too big for SRAM, so FlashAttention tiles it and never materializes it.
- Long context is expensive: every token added makes the matrix quadratically bigger. 128K context is a 128K × 128K matrix per head.
Attention is a cocktail party where everyone whispers to everyone. The more people at the party, the more whispers, and the whispers are the memory (KV cache) and the compute (prefill).
The inference shape
Here's the crucial asymmetry, the one that defines everything: prefill computes the whole N×N matrix once; decode computes only one row of it per token. That's why prefill is compute-bound and decode is memory-bound, and why every optimization in this journey is either making prefill cheaper or making decode faster.
Attention is O(n²) and that one fact explains the entire shape of inference: the KV cache, the roofline, FlashAttention, and why long context is expensive.
Next: the KV cache, where memory goes.
Sources
- Attention Is All You Need: the paper the whole shape of this comes from.
- Fast transformer decoding: one write-head is all you need: multi-query attention.