The KV cache, without the mystery.

When a causal language model generates another token, the earlier tokens have not changed. With consistent positional handling and inference settings, their keys and values do not need to be recomputed. A KV cache saves them for the next decoding step.

This note describes ordinary causal self-attention and a deliberately small cache. Real inference engines also manage memory allocation, batching, and varying sequence lengths.

Prefill and decode

During prefill, the model processes the prompt and stores each layer’s key and value tensors. During decode, the new token produces a new query, key, and value. Its query attends to the cached prefix plus the new key/value pair.

FIG. 005 — Each generated token appends one key and one value per head at each layer. Earlier entries are reused.

There is usually no need to retain past queries: they were used to compute past outputs. The next output needs the new query and all visible keys and values.

A minimal single-head cache

This example receives already projected Q, K, and V vectors. The caller is responsible for positional encoding. Appending with torch.cat is easy to understand but repeatedly reallocates memory; production implementations usually preallocate or page their storage.

python
import torch


class KVCache:
    def __init__(self) -> None:
        self.k: torch.Tensor | None = None
        self.v: torch.Tensor | None = None

    @torch.no_grad()
    def step(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        # One token per call: q, k, v have shape (1, dim).
        self.k = k if self.k is None else torch.cat([self.k, k], dim=0)
        self.v = v if self.v is None else torch.cat([self.v, v], dim=0)
        scores = (q @ self.k.T) / q.shape[-1] ** 0.5
        return scores.softmax(dim=-1) @ self.v


torch.manual_seed(3)
q, k, v = [torch.randn(5, 8) for _ in range(3)]
cache = KVCache()
incremental = torch.cat([
    cache.step(q[t:t+1], k[t:t+1], v[t:t+1])
    for t in range(5)
])

# Compare against all tokens at once with a causal mask.
scores = (q @ k.T) / 8 ** 0.5
future = torch.ones(5, 5, dtype=torch.bool).triu(1)
full = scores.masked_fill(future, float("-inf")).softmax(-1) @ v
torch.testing.assert_close(incremental, full)
assert cache.k is not None
print(cache.k.shape)  # torch.Size([5, 8])

Open complete file

The one-token step does not need a triangular mask: every cached position is in the current token’s past or present. If decoding multiple new tokens at once, the new chunk needs correct causal masking.

Where the memory goes

For standard KV storage, an approximate payload size is:

$$ \text{bytes} = 2 \cdot L \cdot B \cdot T \cdot H_{kv} \cdot D_h \cdot s. $$

READ THE EQUATION

Multiply every cache axis and the bytes per value, then count both the key and value arrays.

Term	What it is and does	What it controls
bytes	Estimated storage for cached K and V values.	The quantity being counted; other memory uses are excluded.
2	One array of keys and one of values.	Doubles the count for both cached projections.
L, B, T	Layer count, batch size, and number of cached positions per sequence.	Each increases cache storage linearly at fixed other dimensions.
$H_{kv},D_h$	Number of KV heads and features per head.	The stored width at each layer, sequence, and position.
s, multiplication dots	Bytes per element, multiplied by the number of elements.	Converts the element count to storage; changing precision changes s.

Check: With L=2,B=1,T=10,Hkv=2,Dh=4,s=2, the cache holds 640 bytes. Doubling cached length gives 1,280 bytes. This estimate omits padding, allocator overhead, model weights, and other activations.

L is layer count, B batch size, T cached length, $H_{kv}$ the number of KV heads, $D_h$ head dimension, and s bytes per element. The factor of two accounts for both keys and values.

For 32 layers, a batch of one, 4,096 tokens, 8 KV heads, a head dimension of 128, and 2-byte elements, that is 512 MiB. This excludes model weights, activations, allocation overhead, and other runtime memory.

What caching does not remove

The new query still attends over the visible prefix, so ordinary attention’s per-token work grows with cached length. The cache avoids recomputing the old representations; it does not make attention over an arbitrarily long prefix free.

Prefill and decode

A minimal single-head cache

Where the memory goes

What caching does not remove

Further reading