How attention works in an LLM, with real numbers
Query, key, value and softmax, measured on a real model: how a token looks back at earlier ones, and how one head lets a model copy what it saw before.
Part 2 of 6 · How an LLM runs
- How an LLM works: from your text to the next token
- How attention works in an LLM, with real numbers
- What the KV cache is, and why LLMs need it
- Prompt caching explained: why an LLM re-reads your whole chat
- What an inference engine does: batching, KV memory and quantization
- Why LLMs run on GPUs: reading a prompt vs writing a reply
In part 1, every layer of a model ran two steps on every token: attention, then a feed-forward network. The feed-forward step works on one token at a time. Attention is the only step where tokens exchange information. This part shows how it works, with numbers from a real model, SmolLM2-135M.
To get at the numbers I rewrote the model in NumPy, reading the same weights file, and checked it against the official runtime. Details are at the end.
Query, key and value
When attention runs, each token's hidden state is multiplied by three weight matrices that belong to that layer. That makes three new vectors per token:
- a query: what this token is looking for,
- a key: what this token can be found by,
- a value: what this token hands over when it is found.
The weight matrices are fixed. The queries, keys and values change with your text, because they're made from the hidden states. In SmolLM2 each one has 64 numbers per head (heads are further down).
Three steps
For one token, attention does this:
- Score. Multiply the token's query with the key of every token up to and including itself, number by number, and add up the products (a dot product). Then divide by 8, the square root of 64. The authors of the original Transformer paper divided because large dot products push softmax to where its gradients get extremely small, which slows training.
- Share. Softmax turns the scores into shares that add up to 100%.
- Mix. Add up the values of those tokens, each weighted by its share. The mix goes through one more weight matrix and is added to the token's hidden state.
A token's key decides how much attention it gets from a query. Its value decides what it hands over.
One real example
I gave the model a list of eight unrelated words, then the same list again:
Words: orange piano seven cloud ladder violet engine marble. orange piano seven cloud ladder violet engine marble
Here is one attention head (layer 17, head 4) at the second "seven". Its query is scored against the keys of all 14 tokens up to there:
| Token | Query · key | ÷ 8 | Share |
|---|---|---|---|
| Words | 39.35 | 4.92 | 2.3% |
| : | 14.33 | 1.79 | 0.1% |
| orange | 29.10 | 3.64 | 0.6% |
| piano | 7.16 | 0.89 | <0.1% |
| seven | −37.74 | −4.72 | 0% |
| cloud | 69.44 | 8.68 | 96.9% |
| ladder | 10.30 | 1.29 | 0.1% |
| violet | −6.45 | −0.81 | <0.1% |
| engine | −14.57 | −1.82 | 0% |
| marble | −46.34 | −5.79 | 0% |
| . | −56.30 | −7.04 | 0% |
| orange | −10.01 | −1.25 | 0% |
| piano | −0.31 | −0.04 | <0.1% |
| seven (itself) | −52.47 | −6.56 | 0% |
The head puts 96.9% of its attention on the word that followed "seven" last time: "cloud". Softmax is steep. A score of 8.68 against a next best of 4.92 is enough for one token to take almost everything.
Then the model predicts the next token, and gives "cloud" 79.9%.
Heads that copy
Researchers call heads like this one induction heads. They complete a pattern "A B … A" with "B" (Olsson et al., 2022). In this prompt, two heads (layer 17 head 4, and layer 18 head 6) put about 90% of their attention on "the word that came next last time", averaged over the repeated list.
It shows in the predictions. In the first copy of the list, the model gave the actual next word about 0.01% on average: the words are random, so it couldn't know. In the second copy it gave the right next word 69% on average. Nothing in the weights knows this list. The model copied it from earlier in the prompt, through attention.
The mask: never look ahead
A token only ever attends to itself and earlier tokens. Later tokens are masked: their scores are set so that softmax gives them exactly 0%. The Transformer paper adds the mask so that the prediction at each position depends only on the tokens before it.
That's what makes training work: the model predicts every position of a text at once, and no position can peek at the word it's meant to predict. It also means a token's attention never changes once it's computed, because no later token can affect it. Part 3 builds on that.
Many heads, many patterns
Attention doesn't run once per layer. SmolLM2 has 9 heads per layer, each with its own weight matrices, and 30 layers: 270 heads in all. Each head learns its own pattern during training. Three I measured:
- The token before. Layer 16, head 0 puts 89% of its attention on the previous token, across the word-list prompt.
- The first token. On "The trophy didn't fit in the suitcase because it was too big.", 52% of all attention, averaged over every head and position, goes to the first token, "The". 116 of the 270 heads give it more than half. One, layer 18 head 8, gives it 100% (to the nearest 0.01%). This is known as an attention sink (Xiao et al., 2023): heads park attention on the first tokens even when those aren't important.
- "It" and "trophy". Layer 3, head 8 gives 46% of the attention of "it" to "trophy". But averaged over all 270 heads, "it" gives "trophy" 3.1% and "suit" (the first token of "suitcase") 2.4%. One head looking at a word doesn't tell you what the model understood.
Try them. Tap any word to see where it looks:
Why it mattered
Recurrent networks, a common design before Transformers, read text one token at a time, each step waiting for the one before. The Transformer paper points out that attention connects every pair of positions in a constant number of steps, so a token can reach any earlier token directly, and every position can be computed at the same time on a GPU.
The price is pairs. Every token scores itself and every token before it, in every head of every layer. For a 1,000-token prompt that's 500,500 scores per head per layer.
Next, part 3: the KV cache. Keys and values never change once they're made, so a model can save them and skip work.
How I measured
- Model: SmolLM2-135M, base model, Apache 2.0: 30 layers, 9 query heads and 3 key/value heads per layer, 64 numbers per head. Weights from the 32-bit ONNX export
onnx-community/SmolLM2-135M-ONNX(revision d0ae683). - The NumPy version: the Llama architecture written out in NumPy, reading the same weight file, so every query, key, score and share can be printed. Checked against ONNX Runtime on a 14-token prompt: the largest difference in any of the 49,152 next-token scores at any position was 0.0002, and the top token matched at every position.
- Averages: the first-token figures skip the first four positions of the prompt, where there is little else to look at. The "about 90%" for the copying heads is averaged over the 8 repeated words.