KV cache: what can a model safely reuse?

Why process the same prefix again?

After reading a prompt, a causal language model predicts one next token. That token then becomes an input for the following prediction. Re-running the entire growing prefix would repeat work for positions the model has already processed. A KV cache retains the earlier key and value vectors at each attention layer, allowing the next forward pass to process only newly supplied positions. The new queries still attend to the retained context. [1]

This article follows the same two-layer NumPy decoder as the attention article: hidden width 32, four query heads, two KV heads, eight channels per head, and float64 arrays. Its weights are random and fixed. The point is to inspect a concrete inference mechanism and compare calculations, rather than evaluate language quality or GPU performance.

Why old keys and values stay valid

For fixed parameters and positions, the state at position \(t\) depends only on the prefix ending at \(t\). Appending a token cannot change it: the causal mask prevents earlier positions from using that future information. The argument holds through every layer. If an earlier layer’s states are unchanged, its projections and the later layer’s permitted attention inputs remain unchanged too. [2]

Consequently, a layer can retain its old keys and values and calculate only the new ones. For each query head, use its group’s K/V; the formula below suppresses batch and head indices:

There is a separate cache for each layer because the projected states differ by layer. A cache entry is also specific to the prefix, weights, position convention, and attention policy that produced it. Two occurrences of the same word generally cannot share a higher-layer entry merely because their token IDs match.

Why not retain Q? The new output rows use new queries. Old queries would reconstruct old output rows, which are already finished. New attention probabilities are still computed, because a new query compares differently with the cached keys. Caching does not freeze the model’s attention distribution. [1]

Prefill and decoding: the exact timing

Suppose six prompt tokens occupy positions 0–5. Prefill processes those known tokens, fills all layers’ caches, and produces the distribution for token 7. A chosen token is not in the cache until a later forward pass processes it.

Token numbering in the final column starts at one; position IDs start at zero. Keeping those conventions explicit prevents a frequent off-by-one error. Every newly processed token still traverses both complete decoder blocks, including their feed-forward sublayers. Only earlier positions avoid another full traversal.

What is actually stored?

After the six-token prefill, each layer stores K and V with shape \([1,2,6,8]\): batch, KV heads, positions, channels. Queries have four heads, but the persistent cache keeps the compact two-head representation. GQA shares these vectors during attention. Physically duplicating them into four cached heads would waste the intended storage saving. [3]

The demo stores keys after RoPE and values without rotation. When processing position 6, rotate its new query and key with position 6. Keep the cached keys’ existing rotations. Rotating old keys again, or resetting the new position to zero, changes the comparisons. Meta’s reference code similarly applies RoPE before writing keys into its cache. [4]

The cache also records position IDs for explicit masking. K/V arrays dominate storage in a large model, but position metadata, allocation capacity, and temporary arrays must be distinguished when reporting total memory.

The mask must know where the new chunk starts

With six cached tokens and one new token, the new query is at position 6 and all seven keys are allowed. A triangular mask constructed for a local \(1\times7\) matrix can incorrectly allow only column 0. The comparison must use actual query and key positions.

Now append a two-token chunk at positions 6 and 7. The first query can see keys 0–6, and the second can see keys 0–7. Chunked input is useful for checking causality even when a generation loop usually supplies one token at a time.

import numpy as np

past_length, new_length = 6, 2
query_positions = np.arange(past_length, past_length + new_length)
key_positions = np.arange(past_length + new_length)
allowed = key_positions[None, :] <= query_positions[:, None]
# Shape: [2, 8].
# First row allows columns 0..6; second allows columns 0..7.

Line 3 constructs absolute new positions. Line 4 covers cached and new keys. Line 5 broadcasts them into a query-by-key comparison. Padding or a sliding window would require additional constraints; this example uses an unpadded contiguous sequence and full causal attention.

Follow one layer’s cache update

The focused fragment below assumes the current layer has projected only the new input states, giving q_new, k_new, and v_new. It uses the executable example’s apply_rope and grouped_attention helpers. Cache concatenation is on axis 2, the position axis in \([B,H,T,d_h]\).

from tiny_llama import apply_rope, grouped_attention

past_length = past_k.shape[2]
new_length = q_new.shape[2]
positions = np.arange(past_length, past_length + new_length)
q_new = apply_rope(q_new, positions, theta=10000.0)
k_new = apply_rope(k_new, positions, theta=10000.0)

all_k = np.concatenate([past_k, k_new], axis=2)
all_v = np.concatenate([past_v, v_new], axis=2)
key_positions = np.arange(past_length + new_length)
new_heads = grouped_attention(
    q_new, all_k, all_v,
    query_positions=positions, key_positions=key_positions,
)
# Retain all_k and all_v for this layer's next call.
# Merge new_heads, apply the output projection, then continue the block.

The first two lengths separate already processed positions from new work. Only new Q/K receive rotations. Concatenation preserves compact KV heads. The attention helper uses the absolute-position mask and returns outputs for new queries only. Their output projection, residual addition, normalization, and SwiGLU computation then proceed normally. Repeating this process at each layer produces the next-token logits.

concatenate is clear for teaching, but repeatedly copying growing arrays is inefficient. Production implementations can preallocate storage or manage it in blocks. That changes storage management while preserving the logical collection of past keys and values. [5]

Calculate the memory instead of guessing

For uniform layers and full attention, the logical K/V payload is:

The leading 2 counts K and V; \(L\) is layers, \(B\) is batch size, \(T\) is cached positions, and \(b\) is bytes per scalar. For our model:

This is an analytical calculation for the K/V arrays, excluding position metadata, model weights, activations, and allocator overhead. Each additional cached token adds 512 payload bytes across both layers. Using four KV heads would double this payload; changing float64 to a two-byte format would reduce it by a factor of four, before considering numerical differences. The tiny model uses float64 to make numerical comparisons easy to inspect.

Compare cached and full-prefix calculations

Fix the supplied tokens and compare full-prefix and cached logits. With tiny_llama.py in your working directory, this checks one six-token prefill followed by a single new token:

import numpy as np
from tiny_llama import TinyLlama

model = TinyLlama(seed=7)
ids = np.array([[1, 2, 3, 4, 5, 6, 7]])
full = model.forward(ids).logits
prefix = model.forward(ids[:, :6], use_cache=True)
step = model.forward(ids[:, 6:], cache=prefix.cache, use_cache=True)
error = np.max(np.abs(step.logits[:, -1] - full[:, -1]))
print("maximum absolute error:", error)
np.testing.assert_allclose(step.logits[:, -1], full[:, -1],
                           rtol=1e-10, atol=1e-10)

The same model instance ensures identical weights. The prefix call supplies cached context; the next call supplies only the seventh token. Compare matching output positions, since the cached call returns just one row. The full example’s --verify also tests multi-token chunks, which can expose an incorrect offset mask.

These schedules should be equivalent mathematically. Floating-point operation order can produce small differences; the code measures them rather than assuming exact equality. Comparing sampled strings is less diagnostic: different sampling decisions create different later prefixes. Check identical supplied tokens first.

What caching saves, and what remains

At a context length of \(T\), one new query still compares with \(T\) keys and mixes \(T\) values. The attention work for one decode step grows linearly with context length; summing over a growing generated sequence remains quadratic in its length. Prefill still has pairwise attention work across the prompt. Cache storage also grows with context, and reading it can become a serving bottleneck. [5]

Exact reuse requires compatible weights, prefix states, positions, and attention rules. Editing an earlier token invalidates dependent cached states; changing model parameters invalidates the old cache. Quantization, compression, or eviction can change numerical results and require their own evaluation. A bidirectional encoder generally cannot preserve earlier states unchanged when new visible context is appended. Caching is therefore an inference strategy whose validity follows from the model’s dependencies.

Primary references

1. Hugging Face Transformers: How caching works. Per-layer K/V reuse and cache updates.
2. Vaswani et al. (2017), Attention Is All You Need. Causal decoder dependencies.
3. Ainslie et al. (2023), GQA. Compact shared key/value heads.
4. Meta: Llama 3 reference model implementation. Position-aware RoPE, cache writes, and subsequent attention.
5. Kwon et al. (2023), Efficient Memory Management for Large Language Model Serving with PagedAttention. KV-cache memory management during serving.