Forward-deployed / AI Knowledge Graph
Transformersa foundation module
Lesson 05

Generating text and the KV cache

Last reviewed: 2026-10 · Volatility: medium

TL;DR

A transformer writes one token at a time. It reads the whole text so far, picks the next token, adds it, and reads again. Done naively, each step repeats almost all of the last step's work.

The fix is the KV cache. When a token passes through a layer, the model computes its key and value vectors. Those never change for that token. So the model stores them. For each new token it computes one query, one key and one value, and reads the stored keys and values of every earlier token.

The cache has a price. It grows with every token in every request, and it lives in the memory of the accelerator. For a long context or many users, the cache, not the weights, sets how many requests one machine can serve. Prefill, which reads the prompt in one pass, and decode, which makes one token per pass, have different costs, because decode must read every weight to make a single token.

🎯 For the product leader

Why it matters — The cache explains why long prompts raise price, why concurrency drops as context grows, and why the first token is slower than the rest.

What it changes in your decisions — You size serving by cache memory, not by weights alone. You see why shorter prompts and shared prefixes lower cost.

Ask your eng team — "At our target context length, how many requests fit on one accelerator beside the weights, and what limits us: compute or cache memory?"

Risk if ignored — A capacity plan built on the weights alone. The first week of real traffic runs out of memory at a fraction of the planned concurrency.

The mental model: notes you keep versus notes you rewrite

Imagine writing a report one sentence at a time. After each sentence you could re-read the whole page and re-derive your notes on every earlier sentence. Or you could keep the notes and add one line. The cache is the notebook.

The KV cache

Keep the notes, add one line per token

Keys and values of earlier tokens never change, so the model stores them.

No cache: recompute everything
step 1
step 2
step 3
step 4

Every step re-derives the keys and values of every token so far. Work grows with n².

Cache: compute only the new token
step 1
step 2
step 3
step 4

Blue boxes are read from memory. The green box is the one new computation. Work grows with n.

Size: 2 × layers × key-value heads × head size × tokens × bytes · tiny run: 107,520 → 15,360 key and value multiply-adds (7.0× fewer)

7B-class shape at 8,192 tokens, 16-bit: 4.0 GiB per sequence, or 1.0 GiB with 8 key and value heads.

Mechanics

Prefill and decode

Prefill processes all prompt tokens at once. The multiplies are large matrix by matrix products, which hardware handles well. The model reads its weights once and uses them for hundreds of tokens.

Decode produces one token per pass. Each pass multiplies the weights by a single vector. The model must still read every weight from memory to do it. Each 16-bit weight takes 2 bytes and does one multiply-add, which is 2 FLOPs. That is 1 FLOP per byte. Prefill over 512 tokens does about 512 FLOPs per byte. Decode is limited by memory speed. Prefill is limited by compute. The Tensors module explains the ratio.

What the cache holds

For each layer and each token, the cache holds one key vector and one value vector. The size in bytes is

2 × layers × kv_heads × d_head × tokens × bytes per number

The 2 is for keys and values. In a standard model, kv_heads × d_head equals d_model.

What it saves

Without the cache, each step recomputes keys and values for every token so far. With n tokens in total, the work is proportional to n² over the whole generation. With the cache, it is proportional to n.

How it limits serving

The accelerator holds the weights once. Each active request then needs its own cache. After the weights, the leftover memory divided by the cache per request gives the number of requests that fit. Techniques such as paged attention manage this memory without waste. Grouped-query attention, covered in lesson 6, cuts the cache by sharing key and value heads.

Tradeoffs and decisions

  • Memory vs. recompute. Keeping the cache costs memory. Dropping it costs compute. Almost all production systems keep it.
  • Precision. A 16-bit cache is the default. Lower precision saves memory and can cost quality. See quantization.
  • Prefix sharing. If many requests share a prompt prefix, the cache for that prefix can be reused. This is the basis of prompt caching.
  • Batch size vs. latency. More requests per pass use the weights better. Each needs cache memory.

Failure modes

  • Weights-only capacity plan. Memory for the weights fits. Cache memory for real traffic does not.
  • Long prompts at high concurrency. Cache memory per request multiplies by the number of requests.
  • Cache fragmentation. Reserving the maximum length for every request wastes memory.
  • Mixed-up latency. A slow first token, from prefill, is blamed on decode, and the reverse.
  • Stale cache. A cached prefix reused after the prompt template changed.

Under the hood

The file is transformers/code/lesson5_kv_cache.py. The model has a step method that handles one token and a cache of keys and values.

def generate_with_cache(model, prompt, steps):
    cache, tokens = model.new_cache(), list(prompt)
    for pos, t in enumerate(prompt):                      # prefill
        logits = model.step(t, pos, cache)
    for i in range(steps):                                # decode
        tokens.append(argmax(logits))
        if i < steps - 1:
            logits = model.step(tokens[-1], len(tokens) - 1, cache)
    return tokens, cache

The no-cache version calls model.forward on the whole list at every step. The file runs both and compares them.

Worked example: the same text, seven times less key and value work

This example is invented. The numbers come from transformers/code/lesson5_kv_cache.py.

Same output. A 2-layer tiny model reads a 6-token prompt and generates 10 more tokens. With the cache and without it, the 16 tokens are identical.

Less work. Counting the multiply-adds spent on keys and values:

Method Key and value multiply-adds
Recompute everything at each step 107,520
Keep a cache 15,360

That is 7.0 times fewer. The gap widens as generation gets longer.

The cache size matches the formula. The cache holds 960 numbers for 15 tokens. The formula 2 × layers × d_model × tokens gives 2 × 2 × 16 × 15 = 960.

Real sizes, at 16-bit.

Model shape Tokens Cache per sequence
GPT-2 small (12 layers, 12 heads of 64) 1,024 36 MiB
7B-class, 32 layers, 32 heads of 128 8,192 4.0 GiB
Same, with 8 key and value heads 8,192 1.0 GiB

Concurrency on an 80 GiB accelerator. This ignores activations and overhead, so real numbers are lower. The dense 7B-class card has 12.6 GiB of weights and 4.0 GiB of cache per 8,192-token request. That leaves room for 16 requests. The grouped-query card has 11.1 GiB of weights and 1.0 GiB of cache. That fits 68 requests. The same hardware serves four times as many long requests.

Practitioner checklist

  • Do we compute cache memory per request at our target context length?
  • Is the capacity plan weights plus cache, with a margin for overhead?
  • Do we measure time to first token and time per token separately?
  • Are shared prompt prefixes cached and reused?
  • Is the serving stack managing cache memory without fixed reservations?
  • Do we know the key and value head count of each model we compare?

Sources

  • Vaswani et al., Attention Is All You Need, NeurIPS 2017 (arxiv.org/abs/1706.03762): autoregressive decoding with masked attention.
  • Pope et al., Efficiently Scaling Transformer Inference, 2022 (arxiv.org/abs/2211.05102): the cost of the key-value cache and the memory-bound decode phase.
  • Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023 (arxiv.org/abs/2309.06180): cache memory waste in serving and paged allocation.
  • The tiny model, the model shapes and every number are invented or illustrative. They come from transformers/code/lesson5_kv_cache.py. This module's sources were not re-opened during the build, and its fact-check is pending.
Core insight: A transformer is one block repeated, so almost every question about it is a counting question: parameters, multiply-adds and cache.
Part of the learning modules collection.