Home / Chapter 1 · How a model reads and predicts
Last edited · 6 min read
Attention
Attention is the only place in a transformer where tokens exchange information: each token looks at all the earlier ones and takes a weighted mix of their content. The cost of long context comes down to this mechanism.
In plain wordsReading “it was hungry”, your eyes jump back to “cat” to know who is meant. Attention is that glance back, done at every word and towards all the previous ones at once.
Click a token to see which earlier tokens it looks at. The weights in one row always add up to 100%
A row is the token doing the looking, a column is the token being looked at. The hatched half is the future, hidden by the mask. Real weights of one head of the small open model Qwen3-0.6B-Base (layer 10, head 3, counted from 0), with a start token added at the front.
Mechanism: Q, K, V
- Each token projects its vector into three: a query Q (what I’m looking for), a key K (what I can be found by) and a value V (what I pass on). The weights are
softmax(Q·Kᵀ/√d): the dot product measures the match, dividing by the square root of the head dimension keeps the softmax from saturating, and the softmax gives weights that sum to 1. The output is a weighted sum of the Vs. In the example, the Q of “it” matches the K of “cat” best, apart from the start token. - Causal mask: a token sees only itself and earlier tokens. This lets training teach prediction at every position of a text at once, and during generation the K and V of old tokens never change, so they can be kept in the KV cache (see “The generation loop and KV cache”).
- Multiple heads: in Llama 3 70B the 8192-number vector is projected into 64 query heads of 128 dimensions (and 8 key/value heads). Each head sees the whole vector through its own projection and does its own lookup, and an output matrix combines the results. A few heads can be named, e.g. one that copies a pattern already seen in the text; most can’t.
- The attention operation itself ignores order: without the causal mask and without positional information, “2 − 1” and “1 − 2” would look the same. The mask alone lets a model infer positions (a token can count how many predecessors it sees), and Llama 4 has some layers with no positional encoding at all. Most models still add RoPE: it rotates pairs of Q and K coordinates by an angle proportional to position, so the Q·K product depends on the distance between tokens. Context is usually extended after training by rescaling RoPE and a short training run on long texts.
Lengthen the context and compare three numbers: pairs grow with the square, the KV cache linearly, and attention takes over most of the compute only with a very long prompt
Estimate for Llama 3 70B (80 layers, 64 query heads, 8 KV heads, head dimension 128). Masked attention costs about 2 × n × 8192 operations per token per layer, the rest of the layer about 1.7 billion. The KV cache in BF16 is 0.33 MB per token. Llama 3’s window ends at 128k tokens, so 1M is the same formula beyond its reach.
What it costs
- Prefill computes n²/2 pairs in every head of every layer. The rest of the layer, the projections and the MLP, grows linearly, though, and dominates at short context: in Llama 3 70B attention only catches up at about 100k tokens. A prompt ten times longer therefore costs 10 to 100 times more compute, depending on its length.
- During generation, a new token compares one Q against all the Ks in the cache, so the work per token grows linearly. Memory is what hurts: the KV cache grows linearly with length (43 GB for 128k tokens in Llama 3 70B), and it limits how many conversations fit on a card.
- FlashAttention computes exactly the same result, but in tiles in fast on-chip SRAM, without writing the n × n matrix to GPU memory. Memory becomes linear and the computation several times faster, even though the number of operations doesn’t drop (in training it even rises, because the backward pass recomputes attention).
- The KV cache is shrunk through architecture: GQA shares one K, V pair across a group of query heads (Llama 3 70B: 8 instead of 64), DeepSeek’s MLA stores one compressed vector, and a sliding window limits some layers to the most recent tokens (Gemma 3: 1,024, gpt-oss: 128).
Limits of long context
- The softmax always hands out 100% of the attention, so a head can’t “look at nothing”. Models learn to dump the excess on the first tokens (attention sinks, Xiao et al., 2023); in the widget above, that is the start token. That is why cutting off the start of the context, e.g. in a naive sliding window, breaks the model, even when the start carried no important content.
- Fitting text in the window doesn’t mean the model will use it well. Models use information from the middle of a long context worse than from the beginning and the end (Liu et al., “Lost in the Middle”, 2023). What to put in the context and where is covered in “Context engineering and memory”.
Check yourself
How does attention work, and why is long context expensive?
Each token projects its vector into a query, a key and a value. The weights are a softmax over Q·K dot products divided by √d, and the output is a weighted sum of the values. A causal mask hides the future, and many heads do this in parallel. In prefill the number of pairs grows quadratically with length, but in a 70B model attention only dominates compute at around 100k tokens. During generation memory hurts more: the KV cache grows linearly, about 0.33 MB per token in Llama 3 70B (BF16), and limits the batch. FlashAttention, GQA or MLA, and sliding windows help.
Po polsku
Każdy token rzutuje swój wektor na zapytanie, klucz i wartość. Wagi to softmax z iloczynów Q·K podzielonych przez √d, a wynik to ważona suma wartości. Maska przyczynowa zasłania przyszłość, a wiele głów robi to równolegle. W prefillu liczba par rośnie z kwadratem długości, ale w modelu 70B zaczyna dominować obliczenia dopiero przy ok. 100 tys. tokenów. Przy generowaniu bardziej boli pamięć: KV cache rośnie liniowo, ok. 0,33 MB na token w Llama 3 70B (BF16), i ogranicza batch. Pomagają FlashAttention, GQA lub MLA i okno przesuwne.
Follow-up questions (4)
- Why divide by √d?
- The dot product of vectors with random components has a variance that grows with the dimension. Without scaling, the logits get large, the softmax becomes almost one-hot and the gradients vanish.
- How do MHA, MQA and GQA differ?
- In the number of key and value heads. MHA has as many as there are query heads, MQA one shared head, GQA a few groups (Llama 3 70B: 8 for 64 query heads, so an 8 times smaller KV cache). DeepSeek’s MLA stores one compressed vector instead of K and V. Less KV means bigger batches and longer context. MQA and GQA pay for it with a small loss in quality, while DeepSeek reports MLA matching or beating full MHA.
- What does FlashAttention do?
- It computes exactly the same attention, but in tiles that fit in the GPU’s fast memory, without writing out the full n × n matrix. Same result, linear instead of quadratic memory and fewer transfers, so it runs faster. It doesn’t reduce the number of operations.
- How do models handle a million tokens?
- They combine full attention with cheaper variants, or replace it with them: sliding windows; sparse attention, where each token looks only at the top-k earlier tokens picked by a small, fast scorer (DeepSeek V4 pairs it with a compressed KV cache to serve 1M tokens); linear attention; or Mamba-style layers with a fixed-size state. Add GQA or MLA and rescaled RoPE. Fitting a million tokens is not the same as using them well, so measure long-context quality on your own task.