KV Cache: Reuse Past Keys and Values

A causal transformer predicts the next token from the prefix it has already seen. When a new token arrives, earlier tokens' keys and values remain usable: future tokens cannot change their causal hidden states. A KV cache stores those tensors at every attention layer so decoding can process only the new input token.

Caching saves repeated projections and earlier-token computation. The new query still attends to the growing prefix, so cached dense attention is not constant-time decoding.

Follow one decode step

Step from prompt prefill to the next input token. Compare the projection rows in both lanes, then follow the newest query across every key. Change the K/V-head count to see storage shrink while the query's attention work stays the same.

The words are a fixed teaching sequence, not output from a running language model. The figure counts token rows and score pairs rather than presenting animation timing as inference speed.

Prefill builds; decode appends

Prefill processes the prompt's tokens together, with a causal mask. Each layer produces its prompt K and V tensors and stores them. The final prompt state predicts the first output token.

Decode feeds that output token back into the model. At each layer it computes the new token's Q, K and V, appends K/V at the same sequence position, and lets Q read all cached K/V through that position. The resulting state predicts the following token.

This reuse assumes causal attention, fixed model weights, unchanged earlier position assignments and evaluation-time behavior. It is not a general claim that representations in a bidirectional encoder remain unchanged when its input grows.

What caching saves, and what it still costs

Let T be the prefix length including the new input, D the model width, Hq the query-head count, Hkv the K/V-head count and d the head dimension. Hold model dimensions fixed when comparing growth in T.

For K and V projections at one layer, counting multiply-accumulates:

That new-token projection work is independent of T. But the newest query has a score row of length T in every query head, and its value-weighted sum also visits that prefix:

With L layers, attention therefore grows as O(L Hq T d) per decode step. Over n generated tokens after a fixed prompt, dense cached attention still contributes O(n²) total work. A naive uncached forward over the whole growing prefix repeats earlier query rows too; it is more expensive than the newest-query comparison shown in the figure.

Actual latency also depends on memory bandwidth, kernel efficiency, batch size and hardware. A token-row reduction is not a measured speedup.

Keep the sequence position outside the layer update

A token has one position shared by every layer. Incrementing a single cursor inside each layer's cache update would place the same token at a different position in each layer.

This minimal preallocated example instead receives start_pos from the caller. It accepts both a prompt chunk and a single decode token. K/V tensors use [batch, Hkv, new_tokens, head_dim].

# Teaching example: explicit position shared by all layers.
import torch


class KVCache:
    """Store one compact K/V tensor per layer."""

    def __init__(self, layers, batch, kv_heads, capacity, head_dim):
        shape = (layers, batch, kv_heads, capacity, head_dim)
        self.keys = torch.empty(shape)
        self.values = torch.empty(shape)

    def update(self, layer_idx, new_k, new_v, start_pos):
        """Write a chunk; return this layer's prefix through the chunk."""
        assert new_k.shape == new_v.shape
        end_pos = start_pos + new_k.shape[2]
        self.keys[layer_idx, :, :, start_pos:end_pos, :] = new_k
        self.values[layer_idx, :, :, start_pos:end_pos, :] = new_v
        return (
            self.keys[layer_idx, :, :, :end_pos, :],
            self.values[layer_idx, :, :, :end_pos, :],
        )

For prefill, every layer receives start_pos = 0 and a prompt-length chunk. For the next decode step, every layer receives start_pos = prompt_length and one token. The caller advances the sequence position after the complete model forward, not after each layer. Production code also manages device, dtype, capacity and request-specific positions; the example uses float32 CPU tensors for clarity.

The original Llama implementation follows this ownership: the model passes the same start_pos to every layer, and each layer writes its own cache slice. The Transformers cache explanation describes the corresponding prefill and iterative-decode flow.

Calculate the memory trade-off

For B sequences, L layers, T cached positions, Hkv stored K/V heads, head dimension d, and b bytes per element:

The leading 2 counts both keys and values. Use K/V heads, which can be fewer than query heads under GQA or MQA.

For Llama 2 70B, B = 1, L = 80, T = 4,096, Hkv = 8, d = 128, and bf16 (b = 2) give 1.25 GiB of raw cache payload. A hypothetical MHA version with 64 K/V heads and all other dimensions equal would use 10 GiB. The GQA page records the head configuration and its sources.

These are tensor sizes, not total device-memory estimates. Model weights, temporary workspaces, allocator rounding, block metadata and concurrent requests add to observed memory.

Shrinking and managing the cache

Each technique changes a different part of that ledger:

• MQA/GQA stores fewer K/V heads while preserving independent query-head score work. The checkpoint must support that architecture.
• Sliding-window attention bounds retained positions, at the cost of discarding direct access to older tokens. It is an attention-pattern choice, not a free replacement for full attention.
• Cache quantization reduces bytes per element. Scales, packing and quality effects require separate accounting.
• Paged caches manage allocation and permit shared prefix blocks. Paging reduces allocation waste; it does not eliminate the payload or guarantee zero fragmentation.

The score and value calculation is covered in scaled dot-product attention; sliding window attention adds a different restriction on which past positions remain directly visible.

Related concepts