Inference and the KV cache

Prefill, then decode

With a cache, generation runs in two separate phases. A restaurant kitchen is a good picture. First the cooks do all the chopping and prep in one go. Then they plate dishes one at a time.

1. Prefill. The model reads the whole prompt in one pass and fills the cache. It handles all the prompt's tokens at the same time, in parallel. This is a big block of matrix math, the kind of work a GPU does efficiently. How long it takes is the time to first token (TTFT). A longer prompt means a longer wait.
2. Decode. After that, the model makes one token per step. Each step does little arithmetic. But it must read all the model's weights, plus the whole cache, from the GPU's memory. So the limit isn't arithmetic. It's memory bandwidth: how many bytes per second the GPU can read from its memory. The speed you see is tokens per second.

This difference is one reason APIs charge more for output tokens than input tokens. Input tokens are processed in bulk during prefill. Each output token costs a full step of its own, one after another.

The cache is big

Now count what the cache holds. You know these words from earlier lessons:

• A model is a stack of layers, and each one has its own attention.
• Each layer's attention is split into several heads.
• Each head's key (or value) is a list of numbers. The length of that list is the head size.
• Each number takes up some bytes of memory. Models usually store 16-bit numbers, which take 2 bytes each.

Many models let several heads share one set of keys and values (more on that below). So the count uses KV heads: the heads that keep their own keys and values.

For every token, every layer stores one key and one value for each KV head. Multiply it all out:

The 2 at the front is there because each token stores two things: a key and a value.

In plain words: each key or value is a list of numbers. Count how many numbers one token adds across all layers and KV heads, and double it for keys and values. Then multiply by the number of tokens and by the bytes each number takes.

layers, kv_heads, head_size, bytes_per_number = 2, 1, 4, 2
per_token = 2 * layers * kv_heads * head_size * bytes_per_number   # 2: keys and values
print("bytes per token:", per_token)
print("bytes for 3 tokens:", per_token * 3)

Real models

Real models are far bigger. The code below works out the cache for two Llama 3 models. It also tries a made-up version of the 70B model that gives every query head its own KV head.

The sizes print in KiB and GiB. A KiB is 1,024 bytes. A GiB is 1,024 × 1,024 × 1,024 bytes, a little over a billion. "128K tokens" means 128 × 1,024 = 131,072 tokens.

The last two lines ask two practical questions about the 8B model on one GPU with 80 GB of memory. How many conversations fit? And how fast can it decode?

def kv_cache_bytes(layers, kv_heads, head_dim, tokens, batch=1, bytes_per_value=2):
    return 2 * layers * kv_heads * head_dim * tokens * batch * bytes_per_value   # 2: keys and values

GiB = 1024 ** 3
models = [
    ("Llama 3 8B", 32, 8, 128),
    ("Llama 3 70B", 80, 8, 128),
    ("70B with one KV head per query head", 80, 64, 128),
]
for name, layers, kv_heads, head_dim in models:
    per_token = kv_cache_bytes(layers, kv_heads, head_dim, 1)
    full = kv_cache_bytes(layers, kv_heads, head_dim, 128 * 1024)
    print(f"{name:<36} {per_token / 1024:>6,.0f} KiB per token   {full / GiB:>6,.0f} GiB at 128K tokens")

weights = 8e9 * 2                      # Llama 3 8B in 16-bit: 2 bytes per parameter
free = 80e9 - weights                  # what's left on an 80 GB GPU
per_conversation = kv_cache_bytes(32, 8, 128, 4_000)
print(f"\nconversations of 4,000 tokens that fit beside the 8B model on one 80 GB GPU: {int(free // per_conversation)}")
print(f"decode speed limit for one conversation at 3.35 TB/s: {3.35e12 / weights:.0f} tokens per second")

The 8B model adds 128 KiB per token, which comes to 16 GiB for a 128K-token conversation. The 70B model adds 320 KiB per token, or 40 GiB at 128K tokens. Three things stand out.

Sharing keys and values

Picture a study group of 8 people sharing one textbook instead of buying 8 copies. Llama 3 70B does the same with its heads. It has 64 query heads but only 8 KV heads. Each KV head is shared by a group of 8 query heads. This is called grouped-query attention (GQA).

Sharing cuts the cache 8 times, from 320 GiB to 40 GiB for a 128K-token conversation. Quality drops only a little. Multi-query attention goes all the way, with a single KV head shared by every query head.

Memory decides how many users a GPU can serve

The 8B model's weights take 16 GB: 8 billion numbers at 2 bytes each. That leaves 64 GB free on an 80 GB GPU. Each 4,000-token conversation needs about 0.52 GB of cache. So 122 conversations like that fit at once, as the code printed.

Everyone shares the weights, but every conversation needs its own cache. Servers run many conversations together in a batch. Each read of the weights then makes one token for all of them, like a bus that carries many passengers on one trip. Phase 6 covers this in more detail.

Decode speed is set by memory

Making one token means reading every weight once. That gives a speed limit:

The sign $\le$ means "is at most".

In plain words: each new token needs one full read of the weights. Divide how fast the GPU reads by how much there is to read. That's the most tokens it can make in a second.

Try small numbers. A model has 10 GB of weights, and the GPU reads 100 GB per second. It can read the weights 100 ÷ 10 = 10 times a second, so it makes at most 10 tokens per second.

weight_bytes = 10e9        # 10 GB of weights
bandwidth = 100e9          # the GPU reads 100 GB per second
print(bandwidth / weight_bytes, "tokens per second at most")

Now the real case. An H100, a common data-center GPU, reads 3.35 TB per second (3.35 trillion bytes). Dividing by the 8B model's 16 GB gives about 209, as the code printed. So even this GPU tops out at around 200 tokens per second for a single 8B conversation. That's a simplified limit: it leaves out reading the cache, so real speeds are lower. Storing each weight in fewer bytes (quantization, Phase 6) raises the limit.

Prompt caching is a KV cache you keep

You met prompt caching in Phase 1. It's the same idea, carried across requests.

Many requests start with the same long text: a system prompt, a product catalog, a contract. That shared start is the prefix. The provider can keep the KV cache for the prefix. Next time, it skips the prefill for that part.

That's why cached input tokens are billed at a fraction of the normal price. It's also why the first token arrives sooner.

There's a strict rule, though. The cached part must be an exact, unchanged prefix. Remember that each key and value depends on its own token and every token before it. Change one early token, and every key and value after it changes too. Think of a tower of blocks: you can keep the bottom of the tower only if every block in it is exactly the same.

For marketers: put the long, stable material at the start of the prompt. That means brand guidelines, the catalog and examples. Put the part that changes per request, such as the customer's question, at the end. Then the stable prefix can be cached.