KV Caching: Why the Prefix Can Stay Put
Prove the reuse layer by layer
At the input, an earlier token's embedding and position are unchanged when later tokens are appended. In a causal attention layer it reads only its prefix, so its output is unchanged. Position-wise normalization and MLPs preserve that independence. Induction over layers shows that earlier hidden states, and therefore their K/V projections, remain unchanged.
This argument assumes evaluation mode, fixed positions, unchanged weights, and a truly causal computation. Full bidirectional attention does not satisfy it: appending or changing a later token can alter earlier hidden states. Cache validity is a dependency property, not just a check that the earlier token IDs stayed the same.
Prefill populates the cache
For a prompt of $P$ tokens, prefill computes all prompt states under a causal mask. Each layer stores its own keys and values. The final logit at the last prompt position supplies the first next-token distribution.
The cache is not a single table of token embeddings. A layer's keys and values are projections of that layer's contextual input. Different layers generally have different caches even for the same token positions.
A decoding step appends one new entry
- Take the previously sampled token as the next input and assign its logical position.
- At each layer, compute the new token's query, key, and value.
- Append the new K/V entry to that layer's cache.
- Read the cache with the new query and finish that token's block computation.
- Use the final vocabulary logits to sample or select the following token.
In this inclusive causal convention, appending before the read lets the query use its own current token. A fused implementation can implement the same dependency without a literal append operation. The offset mask handles multiple new tokens verified at once.
Caching reduces recomputation. The new query still reads a growing history, so decode attention cost depends on prefix length. Enable JavaScript to change the example inputs; the complete calculation remains in the article.
What work remains per new token?
Without caching, a naïve generator repeatedly runs the whole growing prefix. Caching reuses old projections and old layer outputs, so only one new token passes through the MLPs and query/key/value projections. Its attention still scores allowed historical keys and combines their values.
For a fixed-width layer and current length $n$, attention's new-token history read is $O(nd)$ across standard heads, while projection/MLP work is roughly $O(d^2)$. Over $T$ generated tokens, those growing reads can sum quadratically in $T$, even though each step avoids recomputing all previous tokens through every layer.
A concrete memory budget
For equal key/value head width $d_h$, $H_{kv}$ cached heads, $L$ layers, batch $B$, length $n$, and $b$ bytes per scalar, storage is $2BLnH_{kv}d_hb$. The head-sharing chapter derives this count.
With 32 layers, 8 K/V heads, width 128, and two-byte storage, each additional token costs 131,072 bytes per sequence across the model: 128 KiB. A batch of 16 adds 2 MiB of cache per generated token. At long lengths, this dynamic memory can become a major resource even when model weights fit comfortably.
This formula excludes page-table overhead, quantization scales, allocator rounding, temporary workspaces, and recurrent states in hybrid models. A cache-size calculator should label its assumptions rather than report a single “GPU memory required” number.
Stored history must also be read
At batch one, a decoding step can involve relatively little arithmetic compared with the amount of model and cache data transferred. Batching can amortize weight reads, but each request still has its own history. Grouped heads, cache quantization, and optimized layouts can reduce the cost of that history access.
A smaller cache is not automatically a proportional latency reduction: weight bandwidth, matrix efficiency, synchronization, and other layers may dominate. The complexity and prefill/decode explorer separates these resource counts.
When prefix caching across requests is valid
Identical text is not always sufficient. Reuse requires matching token IDs, positions, relevant model state, attention boundaries, and model parameters. Different adapters or prompt templates can invalidate a match. Partial pages may need copy-on-write when requests branch into different continuations.
PagedAttention addresses how logical token histories map to reusable physical cache blocks. It does not change the causal proof of which hidden states may be reused.
Compare equivalent execution paths
In evaluation mode, compute a sequence's attention once with all inputs present and once by growing its cache. Match logits or hidden outputs at every position. Then perturb future inputs and verify earlier outputs do not change. These tests catch errors in append order, masks, and position indexing.
Try it: Why are previous queries usually not kept in the ordinary decode cache?
Only the new query is needed to compute the new token’s attention output. Historical keys and values are the source it reads. Previous query results and previous layer outputs do not need to be recomputed for that step.
Manage the cache and the workload
Continue to cache paging or draft-and-verify generation. The cache role and head-sharing motivation are also discussed in Shazeer's fast-decoding paper.