Inference: sampling and the KV cache

Why does a model give different answers to the same prompt, and why is its first word the slowest?

Training is over and the weights are frozen. Yet once you start using the model, two things look odd.

• Same question, different answers. Ask a chatbot “What is a cat?” twice and you get two answers. But the model is a fixed function: same weights, same input. Where does the difference come from?
• A pause, then a stream. Paste in a long document. There is a wait, then words stream out at a steady pace. Why is the first word slow and the rest fast?

Pause and predict Same weights, same prompt, yet a different answer each time. Where does the difference come from?

The problem

Inference uses the same network, very differently:

Pause and predict The network hands back about 50,000 scores. Suppose you always take the highest one. What goes wrong?

Pause and predict The generation loop from Build GPT has produced 1,000 tokens. To produce token 1,001, how many tokens go through every layer?

A mental model

Half A: choosing a token

After softmax, the scores are a probability for every token. Sampling means rolling a die with those probabilities painted on its faces. “mat” at 40% wins four rolls in ten, not every roll.

That roll is why the same prompt gives different answers. Nothing has broken when it happens: the model produced the same probabilities both times, and the die landed differently.

Three knobs reshape the die. The first is temperature: divide every logit by T before softmax.

Pause and predict Logits [2, 1, 0] give the favourite 67% at T = 1. At T = 0.5, does the favourite’s share rise, fall or stay the same?

Top-k and top-p

Two ways to cut off the tail. Top-k: keep the k most likely tokens. Top-p: keep the smallest group of top tokens whose probabilities add up to at least p. The rest get probability 0; the survivors are rescaled to sum to 1.

For example: Probabilities 50%, 20%, 15%, 10%, 5%. Top-k = 2 keeps the first two. Top-p = 0.9 keeps four (50 + 20 + 15 = 85 is not yet 90; adding 10 gets there).

Why cut the tail at all? A real vocabulary has 50,000 or more tokens, almost all nonsense at any given moment, and tens of thousands of tiny probabilities add up to a real chance of drawing one. A bad token is appended to the input, so every later prediction is conditioned on it. That is also why high temperatures fall apart: they hand the tail a large share of the die.

Half B: not repeating work

Now the second puzzle: why is there a pause before the first word, and then a steady stream? Start with what happens without any trick. Look at the generate() loop from Build GPT. For every new character it runs the whole text so far through the model again, even though almost all of that text is the same as on the previous step. With that loop, every word would cost more than the one before, and a long answer would get slower and slower. A chatbot streams at a steady pace because it avoids that repeated work.

For the newest token, attention compares its query with the key of every earlier token and blends their values. So the model needs the keys and values of all earlier tokens, in every layer.

Pause and predict Token 9 arrives. Do token 5’s key and value, in any layer, change?

As a programmer you know what to do with a pure function that keeps being called with the same arguments: memoise it, which means storing the result the first time and returning the stored copy afterwards. Here, store each token’s k and v the first time, per layer, and read them back afterwards. That store is the KV cache.

It is a cache in exactly the sense of caching in System Design: compute once, keep the result, serve every later read from memory. Here it never goes stale, because the past never changes.

With a cache, this is everything one new token costs, in one layer: it makes one q, one k and one v; k and v are appended to the cache; q is compared with every cached k, and the weights blend the cached v.

Why is the query not cached? A token’s query is used once, when that token looks back. Keys and values are read by every future token.

Prefill and decode

Prefill: run the whole prompt through once, all tokens in parallel, to fill the cache and get the first new token. Decode: then produce tokens one at a time, each step processing one token against the cache.

For example: A 3,000-token document plus a question: one big prefill pass (the wait), then maybe 200 small decode steps (the stream).

Let's see the numbers

Temperature by hand

Three tokens, logits [2, 1, 0]. Divide by T, exponentiate, divide by the total.

The ranking never changes, only the contrast: the weakest token goes from 2% to 19% as T goes from 0.5 to 2.

In the playground’s 11-token example the three junk tokens share 0.3% at T = 1 and 7.8% at T = 3: one roll in thirteen is garbage, about eight garbage tokens in a 100-token answer.

The cache by hand

Take Llama 2 7B (2023): 32 layers, 32 attention heads, each with its own key and value, 128 numbers per head, each number stored in 2 bytes. For one token the cache holds a key and a value, in every head of every layer:

2 × 32 layers × 32 heads × 128 × 2 bytes = 524,288 bytes = 0.5 MB per token

Pause and predict A 4,096-token conversation on this model. How big is its cache?

Current models of this size cache much less. Llama 3 8B (2024) also has 32 layers and 32 query heads, but its query heads share 8 key/value heads, a design called grouped-query attention (GQA), so only 8 keys and 8 values are stored per layer: 2 × 32 layers × 8 heads × 128 × 2 bytes = 131,072 bytes = 0.125 MB per token, four times less. The same 4,096-token conversation needs 512 MB, and ten users need 5 GB, well under the model’s 16 GB of weights. Modern LLM architecture explains how the sharing works.

Fitting those caches around real traffic is the core of designing LLM serving in System Design.

The math

Dividing by T before softmax scales the gaps between logits, and softmax only cares about gaps. Dividing the probabilities afterwards does nothing useful (you will debug that mistake below).

Only tokens is not fixed by the model: the cache grows in a straight line with each user’s conversation.

What you should remember

1. The network’s job ends at the logits. Picking a token is a policy applied outside the model: ÷ temperature → softmax → top-k / top-p → renormalise → roll the die. That random draw at every token is why the same prompt gives different answers.
2. Temperature scales the gaps between logits. T → 0 is greedy (deterministic, loops easily). T > 1 feeds the junk tail. Top-k / top-p cut the tail to exactly zero.
3. KV cache = memoisation. Thanks to the causal mask, old tokens’ keys and values never change (while the window is not sliding), so only the new token is computed. Output is identical to the naive loop.
4. Prefill (whole prompt, parallel, the wait) then decode (one token per step, the stream). Speed is bought with memory: 2 × layers × kv_heads × head_dim × tokens × bytes, per conversation.

Where this appears in a real LLM

In this lesson

• 11 hand-picked logits
• A 2-layer random-weight Transformer in NumPy
• A cache of 123 KB at 123 tokens (if stored as float32)
• One sequence, prefill token by token

In a real system

• 50,000 to 200,000 logits per step, from the trained network
• Dozens of layers, fused GPU kernels
• Gigabytes of cache per long conversation, carefully paged in GPU memory
• Many users batched together; prompts prefilled in parallel

A model has no memory between calls. A chat feels continuous because the application sends the whole conversation again every turn. The context window is the maximum length of that input; nothing is kept in it between calls.

We applied temperature first, then top-k, then top-p, as the Hugging Face library does. Not every implementation uses that order, so the advice several providers give is: adjust temperature or top-p, not both at once.