Gen AI

KV Cache

Store every previous token's key and value vectors instead of recomputing them at every generation step, turning per-token generation cost from quadratic back down to linear.

Without a cache, generating the next token recomputes keys and values for every earlier position from scratch; with a cache, only the new token's key and value get computed and appended

Why Does This Exist?

Generating text one token at a time means running the full model over an input that grows by one token at every single step — how LLMs work established that the whole sequence gets reprocessed each time a new token is appended. Taken completely literally, generating token 1,000 would mean recomputing self-attention's keys and values for all 999 prior tokens again, from scratch, even though every one of those tokens' representations is identical to what it was the last time they were computed — nothing about token 500's key vector changes just because token 999 was just generated.

That's an enormous amount of wasted, repeated computation. Recomputing $n$ keys and values at every one of $n$ generation steps costs $O(n^2)$ total work across a full generation, purely from redundant recomputation — on top of whatever attention's own quadratic cost already charges per step. The KV cache is the fix: since a token's key and value vectors never change once computed, store them once and simply reuse them at every later step.

Think of It Like This

Taking notes once instead of re-reading the whole book every page

Imagine summarizing a book chapter by chapter, but with an absurd rule: before writing the summary for chapter 10, you have to reread and re-take notes on chapters 1 through 9 all over again, even though nothing about those earlier chapters changed. That's obviously wasteful — your notes on chapter 3 were correct the first time and stay correct forever, since chapter 3's actual text never changes.

The sane approach is to take notes once per chapter, keep them in a running notebook, and when you reach chapter 10, only take new notes for chapter 10 itself, appending them to what's already recorded. The KV cache is that running notebook: every earlier token's key and value get computed once, stored, and reused at every subsequent step — only the newly generated token needs fresh computation.

How It Actually Works

What gets cached, and why it's safe to cache

At each attention layer, every token in the sequence has a key vector and a value vector, computed by projecting that token's embedding through the layer's $W_K$ and $W_V$ matrices. Because those matrices are fixed (they're trained weights, not something that changes during generation) and because a given token's embedding at a given position never changes once computed, that token's key and value are also fixed forever after they're first computed. The KV cache simply stores every token's key and value vectors, per layer, as generation proceeds — the diagram's right-hand side shows exactly this: three tokens' keys and values sit in the cache untouched, and only the new fourth token's key and value get computed and appended.

The cost this actually removes

Without a cache, generating token $n{+}1$ recomputes $QK^\top$ using all $n{+}1$ tokens' keys, which means recomputing all $n$ old keys from scratch even though they haven't changed — the diagram's left side shows this waste directly. With a cache, generating token $n{+}1$ only needs to compute one new key and value (for the new token), and score the new token's query against the $n$ already-cached keys plus the new one. Per-step cost for a single new token drops from effectively reprocessing the whole sequence to a single new set of projections plus one row of attention scores — and summed across a full generation, the redundant recomputation this removes is exactly the $O(n^2)$ waste described above.

The tradeoff: memory that grows with every token

The KV cache isn't free — it trades compute for memory. Every generated token adds its key and value vectors, for every layer and every attention head, to a running store that never shrinks during a single generation. The total size scales as:

\text{cache size} \approx 2 \times L \times H \times d_h \times n \times b

where $L$ is the number of layers, $H$ the number of heads, $d_h$ the head dimension, $n$ the sequence length so far, and $b$ the bytes per stored value (the factor of 2 covers keys and values separately). This grows linearly in $n$ , which sounds much better than the quadratic cost it replaces — but for long sequences and many concurrent requests, the KV cache is frequently the actual memory bottleneck a serving system runs into, well before compute becomes the limiting factor.

Show Me the Code

A minimal cache that stores keys and values across generation steps, appending only the new token's projections rather than recomputing everything.

1import numpy as np2
3
4class KVCache:5    def __init__(self) -> None:6        self.keys: list[np.ndarray] = []7        self.values: list[np.ndarray] = []8
9    def append(self, k: np.ndarray, v: np.ndarray) -> None:10        self.keys.append(k)      # only ONE new key per step, never recomputing old ones11        self.values.append(v)12
13    def stacked(self) -> tuple[np.ndarray, np.ndarray]:14        return np.stack(self.keys), np.stack(self.values)15
16
17rng = np.random.default_rng(0)18w_k, w_v = rng.normal(scale=0.3, size=(8, 4)), rng.normal(scale=0.3, size=(8, 4))19cache = KVCache()20
21for step, token_embedding in enumerate(rng.normal(size=(4, 8))):     # 4 generation steps22    k, v = token_embedding @ w_k, token_embedding @ w_v                # ONE new k,v per step23    cache.append(k, v)24
25all_keys, all_values = cache.stacked()26print(all_keys.shape)  # -> (4, 4) — one key row per step, none of them recomputed twice

Each loop iteration adds exactly one key and one value to the cache — nothing from a previous iteration is ever touched again, which is the entire mechanism that turns per-step generation cost from recomputing everything into computing only what's new.

Watch Out For

Forgetting the cache has to be invalidated on a changed prefix

The KV cache is only valid for the exact sequence of tokens it was built from — if anything about the earlier tokens changes (a different system prompt, an edited earlier message in a conversation), every cached key and value computed from that now-stale prefix is invalid and has to be recomputed from scratch. Reusing a cache built from a different prefix produces silently wrong attention scores, not an obvious error, because the shapes still match even though the underlying content doesn't.

Treating cache memory as negligible next to model weights

For a single short request, the KV cache is small relative to the model's own weights sitting in memory. At long context lengths, or with many concurrent requests being served at once, the total KV cache memory across all active sequences can rival or exceed the memory the model weights themselves occupy — which is exactly the pressure that techniques like paged attention exist to relieve. Assuming cache memory is a rounding error compared to model size is a common miscalculation when estimating serving capacity.

The Quick Version

A token's key and value vectors never change once computed, so recomputing them at every generation step is pure waste.
The KV cache stores every token's key and value, per layer, and only computes new ones for newly generated tokens.
This turns the redundant-recomputation cost from quadratic across a full generation down to linear.
The cache trades that saved compute for memory that grows linearly with sequence length, across every layer and head.
At scale, KV cache memory is frequently the real bottleneck in serving, not raw compute.