A KV Cache Stores Layer-Specific Keys and Values

Derive why incremental decoding stores keys and values—but not old queries—for every layer and head. Calculate cache axes, scalar count, and float32 payload for the frozen decoder.

At layer $\ell$ , the normalized residual records produce queries, keys, and values. During later decode steps, a new query must compare with every earlier key and mix every earlier value. Those earlier keys and values are therefore stored.

For standard two-head attention at cached length $t$ :

K^{(\ell)},V^{(\ell)}:(B,h,t,d_h).

In the frozen model, each tensor has shape $(1,2,t,2)$ .

Why Previous Queries Are Not Stored

A previous query was used to calculate the output for its own position. The next token supplies a different query. That new query needs earlier keys and values, not earlier queries.

This does not mean old queries have no analytical value. It means they are not required to calculate the next standard causal-attention output.

Every Layer Owns a Different Cache

Layer 1 keys and values come from Layer 1's normalized input. Layer 2 sees the residual result after Layer 1 and therefore produces different projections. Reusing one layer's cache in another layer preserves shape but changes logits.

The verified wrong-layer intervention produces maximum logit error 3.6990, well above the $10^{-5}$ equivalence tolerance.

Cache Validity Depends on More than Shape

A cache is valid only for the same:

model parameters and architecture;
prompt token IDs and their order;
position convention;
layer assignment;
dtype and compatible device behavior;
request or sequence owner.

If any parameter updates, recompute the cache. Do not carry training-time cache state into a later optimizer step.

Q1. Calculate one layer cache shape

For $B=3$ , $h=2$ , cached length $t=5$ , and $d_h=2$ , what are the key and value shapes in one layer?

Your answer

Answer it first, then check.

Hint

Use batch, head, sequence, head-feature order.

Solution

Both tensors have shape

(3,2,5,2)

Research Boundary

Shazeer, Fast Transformer Decoding: One Write-Head is All You Need

That work studies a multi-query variant that shares key/value heads to reduce incremental-decoding memory bandwidth. The current model retains two distinct key/value heads; its cache formula must not silently use the variant.