LM10-T5 — KV Caching, Prefill and Decode

Learning objectives

After this lesson, you should be able to:

• Explain what the KV cache stores and why it does not change the output.
• Count operations with and without the cache and the bytes of the cache.
• Distinguish prefill from decode and state each one’s bottleneck.
• Name the main techniques for managing the cache: grouped-query attention, quantisation, paging, prefix reuse.

The idea: the past never changes

In a causal transformer the keys and values of token t in layer l depend only on tokens 1 … t, because the mask forbids looking forward. When token t + 1 is appended, nothing about token t’s keys and values changes. The past is fixed; only the newest token is new. So after computing the keys and values of a token once, we can store them and reuse them in every later step.

Note what is stored: keys and values, per layer, per key-value head, per token. Queries are not stored, because a token’s query is used only once, in its own step. Attention weights are not stored either; they are recomputed from the new query and the cached keys.

Exactness

Caching is not an approximation. The cached tensors are the same numbers that a full recomputation would produce (up to floating-point rounding differences from the order of operations). Notebook LM10-L3 tests this: logits from an eight-token prefill followed by four cached single-token steps equal the logits of one full pass to about 10⁻⁶.

Two phases

Inference with a cache has two phases with very different characters:

• Prefill: process the whole prompt in one parallel pass, filling the cache. Many tokens are processed at once, so the matrix products are large; the phase is limited by arithmetic. Its duration is the main part of the time to first token.
• Decode: generate one token per step. Each step projects one token, but must read all the weights and the entire cache from memory. The arithmetic is tiny relative to the data moved, so the phase is limited by memory bandwidth. Its speed is the tokens per second the user sees.

The trade

The cache turns repeated computation into stored memory. That is a good trade until memory becomes the scarce resource, which at long contexts and with many concurrent requests it does. Everything else in this lesson is about managing that memory.

Why decode is slow for a reason you can calculate

The mathematics: work saved, memory spent

Operations

Count the multiply-adds of one attention layer of width d (four projections and the two attention products; feed-forward layers also scale with the number of tokens processed, so the saving is similar). At step t, with t tokens in context:

Summed over steps t = 1 … T, without a cache the cost is about 2T²d² + T³d·(2/3); with a cache about 4Td² + T²d. The projections therefore fall from quadratic to linear in T. Worked example (the lab’s case): d = 4, T = 3: without a cache 72 + 160 + 264 = 496; with a cache 72 + 80 + 88 = 240. For d = 128 and T = 1,000 the ratio is about 611.

Memory

For the 7B configuration (L = 32, H_kv = 32, d_h = 128, 16-bit): 524,288 bytes per token. One request with 8,192 tokens: 4.29 GB. Sixteen such requests: 68.7 GB, five times the 14 GB of weights. With 8 key-value heads: 17.2 GB. A single request with a 128,000-token context needs 67 GB with MHA and, with 8 key-value heads and the same formula, 16.8 GB.

The cost of a decode step

A step moves the weights (W_bytes) and the request’s cache (C(n)) from memory. With bandwidth BW:

W_bytes

weights in bytes, for example 14·10⁹.

C(n)

cache bytes for the n tokens of this request.

BW

memory bandwidth, 10¹² bytes/s in the running example.

With 14 GB of weights and 1 TB/s: at 1,000 tokens of context the MHA cache is 0.52 GB and the step takes 14.5 ms (68.9 tokens/s); at 8,192 tokens 18.3 ms (54.7 tokens/s); at 32,000 tokens the cache is 16.8 GB and the step 30.8 ms (32.5 tokens/s). With 8 key-value heads the same steps take 14.1, 15.1 and 18.2 ms. The cache, not the weights, determines how decoding slows down as conversations grow, and grouped-query attention cuts that slowdown sharply.

Prefix reuse and paging

Two prompts that start with the same tokens (a system prompt, a document) have identical caches for that prefix, because the cache of a position depends only on earlier tokens. Systems store and share the prefix cache so a repeated prefix is not prefilled again: this lowers time to first token and cost. Naive allocation reserves a contiguous block for the maximum length of each request and wastes most of it; paged allocation stores the cache in fixed-size blocks allocated on demand, like virtual memory, and allows blocks to be shared between requests (Kwon et al., 2023).

A worked example: count the savings

The code reproduces the operation counts and the decode-step estimates above. It is the same arithmetic as Lab LM10-L1; change the numbers to your own model and budget.

def ops(d, T):
    no_cache = sum(4 * t * d * d + 2 * t * t * d for t in range(1, T + 1))
    with_cache = sum(4 * d * d + 2 * t * d for t in range(1, T + 1))
    return no_cache, with_cache, no_cache / with_cache

print(ops(4, 3))                            # (496, 240, 2.0666666666666664)
print(round(ops(128, 1000)[2]))             # 611

def kv_bytes(layers, nkv, head_dim, tokens, batch=1, bytes_=2):
    return 2 * layers * nkv * head_dim * tokens * batch * bytes_

print(kv_bytes(32, 32, 128, 1))                                 # 524288
print(round(kv_bytes(32, 32, 128, 8192, 16) / 1e9, 1))          # 68.7
print(round(kv_bytes(32, 8, 128, 8192, 16) / 1e9, 1))           # 17.2

def step_ms(cache_bytes, weights=14e9, bandwidth=1e12):
    return (weights + cache_bytes) / bandwidth * 1000

for n in (1000, 8192, 32000):
    mha = step_ms(kv_bytes(32, 32, 128, n))
    gqa = step_ms(kv_bytes(32, 8, 128, n))
    print(n, round(mha, 1), round(1000 / mha, 1), round(gqa, 1))
# 1000 14.5 68.9 14.1
# 8192 18.3 54.7 15.1
# 32000 30.8 32.5 18.2

The pattern to remember: the first line is the saving from not recomputing, the middle lines are the memory spent, and the last block shows memory traffic turning into slower tokens as the context grows.

Managing the cache

Engineering implications and limitations

Test equality with full recomputation. Compare cached decoding with a full forward pass on random inputs. It catches wrong positions, wrong masks and cache-ordering bugs at once; keep it in the test suite.

Plan memory as weights plus cache at the target load. The cache per token times the context times the concurrency is usually the number that decides whether a deployment fits (Lab LM10-L1).

Know which phase you are optimising. Slow first token means prefill (shorten prompts, reuse prefixes, more compute); slow streaming means decode (quantise, smaller cache, more bandwidth, speculative decoding in Lesson LM10-T6).

Store rotated keys. With rotary embeddings (Lesson LM10-T2) rotate each key by its absolute position when it is created and cache the result.

Cap context and output length. Every extra token costs cache memory and decode time for the rest of the request; limits protect the other users.

Free the cache. Release a request’s cache when it finishes or is cancelled; leaked caches are a common cause of slowly growing memory in long-running servers.

For robots. A robot that keeps a long running context (a history of observations and commands) is a long-context workload on small memory. Bound the history, summarise or drop old entries deliberately, and budget the cache for the worst case.

Common misconceptions

Engineering takeaways

1. Keys and values of past tokens never change in a causal model, so they are stored and reused; queries and attention weights are not.
2. Cached decoding is exact; it replaces a quadratic-in-length recomputation of projections by a linear one (496 versus 240 operations in the d = 4, T = 3 example; about 611× for d = 128, T = 1,000).
3. Cache bytes = 2·L·H_kv·d_h·n·B·b: 524,288 bytes per token and 68.7 GB for 16 requests of 8,192 tokens in the 7B example.
4. Prefill is compute-bound and sets time to first token; decode is memory-bound and its step time grows with the cache.
5. Manage the cache with fewer key-value heads, quantisation, paging, prefix reuse and limits on context.

References and further reading

The counts in this lesson come from the code shown and use a simplified operation count (projections and the two attention products of one layer). The bandwidth figure is a round illustrative number, not that of a specific device.

• Vaswani, A. et al. (2017). Attention is all you need. arXiv:1706.03762.
• Shazeer, N. (2019). Fast transformer decoding: One write-head is all you need. arXiv:1911.02150.
• Kwon, W. et al. (2023). Efficient memory management for large language model serving with PagedAttention. arXiv:2309.06180.

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: inference efficiency and architecture trade-offs: combining quantisation, batching and design choices into a budget.

Related content

Nothing below is required to finish this lesson.

• Before
  LM10-T4: SwiGLU, RMSNorm and Mixture of Experts

  Explain gated feed-forward layers, RMS normalization and sparse expert routing.

• Lab
  LM10-L1: KV-Cache Memory and Behaviour

  Compute KV-cache size and count the operations with and without a cache.

• Lab
  LM10-L3: Implement a KV Cache and Rotary Embeddings in PyTorch

  Add rotary embeddings and a KV cache to attention and verify that cached decoding equals full recomputation.