Key-Value Caching

I reused a key-value cache built for one prompt to continue a different one. Nothing raised an error. The next-token scores were off by 11.6 on the median across twenty trials, and by 28.7 in the worst, and the model picked the right next token in only 40 per cent of them. In a second experiment I built the cache with weights that differed from the real ones by 2 per cent in a single layer, as an adapter swap might. The scores were off by 0.04, and the right token came out in 20 trials of 20.

The first failure is loud enough to find on the first test. The second is quiet enough to survive a benchmark, a code review and a demo. Both are the same mistake: a cache is a piece of state, and it is only correct for the exact computation that produced it. Most of this chapter is about building a cache that is right, and the last part is about everything that can make it silently wrong.

The cache is also where the book’s accounting finally changes. Chapter 5 showed that the naive generation loop processes 1,040 positions to produce 32 tokens, where 48 would be enough. Chapter 9 introduced the cache as a black box. Here it is built, verified against the uncached loop, and measured.

Where this sits

Chapter 4 showed that a causal mask makes the output at an early position exactly independent of later tokens: changing the last token left every earlier logit identical, to zero. That is the property that makes a cache legitimate. Chapter 5 computed the waste of the naive loop and left the saving as NOT_RUN. Chapter 8 computed cache sizes and showed that decode attention is a pure bandwidth problem at useful batch sizes. Chapter 9 measured the two phases with the library’s cache.

The Inference Lab piece added here is the cache itself, two designs behind one interface, with a backend that puts it inside the Chapter 6 engine and enforces a rule: a cache may only be reused if it provably holds the state for exactly this prefix. The engine’s backend contract is stateless, “ids in, scores out”, and the cache lives behind it. That is the pattern every later optimisation in the book follows: state hidden behind a stateless interface must validate itself on every call.

For the Pi series, the cache is the mechanism behind several things the agent volumes take for granted: long sessions that do not recompute their history, and the reuse of a shared prefix, such as a system prompt, across requests. Chapters 14 and 16 build on this. The same mechanism is why a changed system prompt, a swapped adapter or a different model version has to invalidate everything downstream.

Why the cache is exact

In attention, position \(i\) produces a query, a key and a value, all functions of the tokens up to \(i\) only (the layer inputs at \(i\) depend on earlier positions through attention, and on nothing later). Once position \(i\)’s key and value have been computed for a given prefix, they are the same at every later step. Only the new position’s query is new, and a query is used once: it asks its question of the keys and values and is never consulted again. That is why queries are not cached.

So at decode step \(t\), for each layer, the new token needs: its own query, key and value (a few small matrix products), and the keys and values of the \(t-1\) earlier positions (read from the cache). The attention score is then a \(1 \times t\) row. Everything Chapter 8 said about the decode shape applies.

The saving, for a request of prompt length \(P\) generating \(n\) tokens:

For GPT-2 with a 17-token prompt and 32 new tokens, that is 1,040 against 48. The engine in this chapter’s last experiment reports exactly those counts, 48 positions per request with the cache, and the ratio of 21.7 is the waste factor that Chapter 5 predicted.

The cache does not make the model linear in context. Each step still reads every weight, runs every layer, and attends over all cached keys. What it removes is the recomputation of the prefix’s projections and feed-forward networks, which is the great majority of a step’s arithmetic.

What is in it, and how big

For each layer, the keys and values of every position, in the shape the attention kernel wants: \([B, h, T, d_h]\) each. The size is

The formula and the measurement agree exactly in all four rows. A full 1,024-token sequence costs 72 MiB in fp32 and 36 MiB in fp16, linear in length and in batch.

For scale: the model’s weights in fp16 are 237 MiB, and a sequence at the model’s full context is 36 MiB. Free memory on this card after the desktop’s share was about 10 GiB, so, ignoring activations, fragmentation and the desktop’s growth, a rough count says 273 full-length sequences’ caches fit next to the fp16 weights. That is arithmetic and not a measurement of a running system. The point is the ratio: in this model, a handful of sequences’ caches cost as much as the weights, which is the Chapter 8 observation that KV traffic overtakes weight traffic at a batch of about seven.

Writing it

The cache needs three changes to Chapter 4’s forward pass.

Positions. The new tokens sit at positions cache.length ... cache.length + t - 1, so the position-embedding rows must be looked up at that offset, not from zero.

The mask. A prefill chunk of \(t\) new tokens attends to \(\text{total} = \text{cache.length} + t\) keys. Query \(i\) is at absolute position \(\text{total} - t + i\) and may see keys up to and including it. The mask is therefore the upper triangle offset by total - t + 1. For a single new token that offset makes the mask empty, because the token sees everything.

Attention reads the cache. Each layer calls cache.update(layer, k, v), which stores the new keys and values and returns all of them, and cache.advance(t) at the end of the pass moves the length forward once, after every layer has been served.

x = wte[ids] + wpe[arange(offset, offset + t)]
for layer in range(n_layer):
    ...
    keys, values = cache.update(layer, k, v)          # [B, h, length + t, d_h]
    scores = q @ keys.transpose(-2, -1) / sqrt(d_h)
    scores = scores.masked_fill(triu(ones(t, total), diagonal=total - t + 1), -inf)
    ...
cache.advance(t)

Two cache designs sit behind update:

• ConcatCache makes a new tensor by concatenation at every step. It is what a straightforward implementation does and what the library’s default cache does in essence. Each step allocates a new, slightly larger tensor and frees the old one.
• StaticCache allocates once, for the maximum length, and writes new keys and values in place. It returns slices of the buffer up to the current length, so values beyond the length are never read, and reset() just sets the length to zero without clearing the buffer.

The second design’s reset is the sort of thing worth a test, and there is one: after a buffer holds one request’s cache, a reset, and a different request, a following step is bit-identical to the same step on a fresh buffer. The stale values are physically present and logically invisible.

It is the same function

The first requirement is that nothing has changed. I decoded 48 tokens greedily from the 17-token prompt on the GPU and CPU with four implementations: no reuse at all, the concat cache, the static cache and the library’s cache.

At the level of scores, I computed at every step the next-token logits two ways: from the cache, and by recomputing the whole sequence. The difference had a median of \(6 \times 10^{-5}\) and a maximum of \(6.7 \times 10^{-4}\) on the GPU (\(5.3 \times 10^{-4}\) on the CPU), against logits whose magnitude is in the tens to hundreds. Those are rounding differences: the cached path computes the same sums in a different order and with different matrix shapes. The disagreement between my uncached forward and the library’s was \(6 \times 10^{-5}\) to \(9 \times 10^{-5}\), the same scale. A cache is not bit-identical to recomputation, and it should not be expected to be. It is equal within the arithmetic’s resolution, and Chapter 39 is where “within tolerance” is turned into a reproducibility policy.

The tests make these the standing checks: cached and uncached greedy decoding agree for both designs on a small randomly initialised model; cached logits match the library model; a chunked prefill matches a one-shot prefill; and the byte formula matches the tensors.

What it saves

The time for one decode step, with a context of \(T\) tokens already present, with and without reuse. Both use the same hand-written forward pass; the “no cache” step processes the new token together with the whole prefix. Two runs:

Speedup of the concat cache over recomputation:

On the CPU the cache is a large win early. A CPU is arithmetic-bound, and avoiding the prefix’s arithmetic matters from the first token: 2.3 times faster at 16 tokens of context, and more than 24 times at 512.

On the GPU the cache is worth nothing for short contexts. At 16 tokens of context the cached and uncached steps are within 20 per cent. The GPU step is launch-bound at about 9 ms (Chapters 7 and 9), and removing arithmetic that was never the bottleneck changes nothing. The cache starts to pay off where the uncached step leaves the launch floor, at a few hundred tokens: 3 to 4 times faster at 512, and 5 to 11 times at 960. The cached step itself stays flat at 8 to 16 ms across the range, which is the same floor again.

The 960 row’s spread (8.4 to 16.5 ms for the concat cache between two runs) is as large as the effect between designs, and I would not read a ranking of concat against static from it.

Inside the engine

The Chapter 6 engine, 32 greedy tokens, with the cache behind its backend. The “uncached” row is the library’s forward pass called on the full sequence at each step.

On the CPU the cache is a 3.5-fold improvement end to end for this 49-token request. On the GPU, my cached backends are slower than the library without a cache: 9.9 to 10.1 ms per token against 8.5. I see three reasons and have isolated none. The hand-written forward pass is unfused eager code and issues more operations than the library’s, so in a launch-bound regime it pays more. At 17 to 48 tokens, the cache saves nothing the GPU would have noticed. And this comparison mixes two implementations, which is why the earlier table, which uses my forward pass on both sides, is the cleaner measurement of what reuse buys.

One honest caveat about absolute numbers. The library’s cached decode step took about 11 ms in Chapter 9 and 5 to 8 ms in this chapter’s runs, on the same machine, with the same model. I checked whether selection and the host readback were the difference; they account for about one millisecond. The rest I put down to machine state, which I did not control. Comparisons within a run are sound, and absolute times from different sessions should not be compared.

What it costs in memory

The cache is a growing allocation, and the two designs treat the allocator differently. I generated up to 960 tokens at batch sizes of 1 and 8 in fp32, recorded the peak extra allocation, and read the allocator’s reserved memory at the end (about 0.5 GiB of every figure below is the model itself):

The peak allocation barely differs, 4 to 9 per cent above the cache itself in all four cases, because the transient doubling in concatenation affects only one layer’s tensors at a time. But the allocator tells a different story. At batch 8 the concat design left 2.9 GiB reserved against 1.1 GiB for the static one, about 1.8 GiB more, to hold a 540 MiB cache. Every step frees a tensor of one size and allocates one a little larger, and the freed blocks cannot be reused for the larger requests. The allocator keeps them as a cached pool, which counts against the GPU. The static design pays 6.7 per cent extra up front (it reserves the full 1,024 positions) and has no growth to fragment.

That is a result about this allocator, this pattern and this one batch size, and I do not know how it scales. It does explain why serious engines preallocate and why Chapter 14 goes further and manages the cache in fixed-size blocks.

State identity

A cache is correct for a specific computation. For GPT-2 that computation is identified by the weights, the tokenisation, and the exact sequence of tokens whose keys and values it holds. The chapter’s backend enforces it on every call. A call to next_logits(ids) is served from the cache only if the cache holds exactly ids[:-k] for the call’s new k tokens. Otherwise the cache is discarded and the prefix recomputed. Identical ids (nothing new to process) are also treated as a miss. In the engine’s experiments every request missed once, on its first call, and then hit for the remaining 31 steps: 5 misses and 155 hits in 160 calls.

To make the identity portable, the module defines prefix_key(manifest_sha256, token_ids): a hash of the artefact manifest from Chapter 2 and the exact token ids. Two requests share a cache only if their keys match. The manifest is part of the key because a cache is a function of the weights, which is the next experiment’s subject.

What goes wrong

Five ways to reuse a cache wrongly, over 20 random trials each. Each trial builds a 24-token prefix, a random next token, and compares the reused cache’s output with a correct computation:

The loud failures are loud. A cache from another prompt or a stale position makes the scores unrecognisable, and the right token appears at about chance. A forgotten reset on the static buffer gives an intermediate failure with the most insidious feature: the next token was still right in 70 per cent of trials, so a casual test with a few prompts would pass.

The last row is the dangerous one. The prefix was computed by weights with one layer’s feed-forward output scaled by 1.02, and continued by the real ones. The scores moved by 0.04 and the argmax was correct 20 times out of 20. For comparison, that error is the same order as an fp16 matrix product (Chapter 7 measured a relative error of \(3.6 \times 10^{-4}\)), and well within the range someone would call “numerically fine” and move on from. Yet the output is not the output of the model that was meant to run: it is a mixture of two models, and no log line will say so.

The remedy is not a cleverer check on the numbers. It is identity. The cache key includes the artefact manifest, and a change to weights, an adapter or the tokeniser changes the key, so the old cache is not found. This is the Chapter 2 manifest earning its place: the thing that distinguishes the two weights is the digest, and the digest costs nothing to compare.

Sharing a prefix

The same identity makes sharing safe. Two requests that begin with the same 512 tokens can share the cache for that prefix. I built the prefix cache once (36 MiB in fp32), forked it into two copies, and continued each with a different token. Both continuations matched an independent computation of the full sequence to \(6 \times 10^{-5}\), and the parent cache’s length and contents were untouched. A real system would not copy: it shares the blocks and copies on write, which is the point of Chapters 14 and 16. The measurement shows that the idea is sound and what it saves, the prefill of 512 tokens, but not how to implement it efficiently.

What the numbers do not show

• GPT-2 small on a launch-bound GPU. The “cache is useless at short context” finding is about a model whose step is dominated by launch overhead. For a larger model, where a step is bytes-bound, the cache would matter from the first token.
• My forward pass is not the library’s. It is unfused and slower on the GPU. The cache comparisons that matter use it on both sides.
• Timings drift. Absolute step times differ by about 2× across sessions, and the two-run ranges above include real noise.
• One batch-size pair for memory. The allocator result is for batch 1 and 8 at 960 tokens.
• Random-data hazards. Effect sizes for stale-prefix and position errors come from random tokens. With real text, the mistakes may be less (or more) obviously wrong.
• One perturbed layer. The “different weights” hazard scaled one layer by 2 per cent. A different perturbation gives a different, not necessarily smaller, error.
• No paging, quantisation or eviction. Those are Chapters 12, 14 and 16.
• The 273-sequence figure is arithmetic. It ignores activations and every other process on the card.

Inference Lab increment

The Lab now has:

• ConcatCache and StaticCache behind update, advance, reset, fork and nbytes, and cache_bytes(...) for the formula.
• forward_cached(ids, weights, config, cache), a generalisation of the Chapter 4 oracle that handles prefill, chunked prefill and decode by the same code.
• CachedBackend, a Chapter 6 backend that holds a cache and validates it on every call, with counters for calls, reuses, resets and positions processed.
• prefix_key(manifest, tokens), the identity of a cache’s contents.

The rule: a cache may only be reused when its key matches, and a change in the artefact is a change in the key. A faster path that cannot say why its state is valid is not an optimisation, it is a bug with a delay.

What you should now be able to do

• Explain why cached keys and values are exact, using the causal mask, and why queries are not cached.
• Compute a cache’s size from model dimensions, and the positions processed with and without it.
• Change a forward pass to read from a cache: position offsets, the offset causal mask, and when to advance the length.
• Verify a cache against recomputation, and say what tolerance to expect.
• Say where a cache helps (arithmetic-bound steps, long contexts) and where it does not (launch-bound steps, short contexts).
• Choose between growing and preallocated caches with the allocator in mind.
• Describe the failure modes of a wrongly reused cache and design the key that prevents them.

Key terms

• KV cache: stored keys and values of earlier positions for every layer.
• Static cache: a preallocated cache written in place.
• Concat cache: a cache grown by concatenation.
• State identity: the exact model, tokens and positions that a cache is valid for.
• Invalidation: discarding a cache because its identity no longer matches.
• Fork: copying a cache so that two continuations can diverge.
• Allocator fragmentation: memory held by the allocator but unusable for the next request.
• Prefix key: a hash of the artefact manifest and the token ids a cache holds.

Reader exercises

1. Remove the total - t + 1 mask offset and decode with the cache. At what step does the output first go wrong, and by how much?
2. Compute the cache size for a 32-layer, 32-head, 128-dimension model at 8,192 tokens in fp16. How many sequences fit next to its weights in 24 GiB?
3. Extend CachedBackend to keep several caches, keyed by prefix_key, and reuse the longest matching prefix. What eviction rule would you use?
4. Perturb a different layer, or the embedding table, by 2 per cent and repeat the “different weights” hazard. Does the next token still match?
5. Measure the allocator’s reserved memory for batch 32 with each design. Does the gap widen?
6. Implement a rollback: cut a cache back to length \(k\). When is this safe, and what must the key become?
7. Time the cached step for a hand-written forward pass with fewer operations, for example by fusing the layer norms. How much of the 9 ms is launch overhead?

References

• Efficiently Scaling Transformer Inference — analysis of the cost of inference, including memory and cache considerations.
• Efficient Memory Management for Large Language Model Serving with PagedAttention — documents the memory waste of contiguous per-sequence caches and the block-based remedy used in Chapter 14.
• Transformers: cache strategies — official documentation of dynamic and static caches in a widely used library.

Evidence in this chapter

Implementation status

• Manuscript: content/books/inference/10-chapter.md
• Runnable examples: examples/ch10/kv_cache.py (original toy), examples/ch10/cached_gpt2.py (the cache, the backend and the measurements)
• Tests: tests/test_ch10_kv_cache.py, tests/test_ch10_cached_gpt2.py (parity of both designs against recomputation and the library, chunked prefill, the byte formula, static-buffer reset, wrong-reuse detection, fork independence, backend hit and miss accounting; all on a tiny model, so no download or GPU is needed)
• Evidence: research/ch10/sources.md, research/ch10/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about paged or quantised caches, larger models or production serving. Chapter 11 returns to the question that keeps surfacing, whether an operation is limited by arithmetic or by memory, and builds the tools to answer it for a whole forward pass.