KV Caching

Why generation does not redo the past at every step: Keys and Values get computed once and kept.

The intuition: the past is settled, so compute it once

During training, the model processes a whole sequence in parallel. During generation it produces one token at a time, feeding each new token back in to produce the next. Done naively, every new step would recompute the Keys and Values for every earlier token, in every layer, repeating almost all of the previous step's work.

But under a causal mask (Part 7), an earlier token's representation never depends on a token that comes after it. Once the model has processed “The cat sat on the,” adding “mat” doesn't change anything about “cat.” So “cat”'s Key and Value, computed once, stay valid for the rest of the generation. Storing them and reusing them is the KV cache.

flowchart TD
  subgraph CACHE["cached from earlier steps"]
    direction LR
    T1["K,V: The"] --> T2["K,V: cat"] --> T3["K,V: sat"] --> T4["K,V: on"] --> T5["K,V: the"]
  end
  NEW["new q,k,v: “mat”
(computed now)"]
  NEW -- "Query compares against every cached Key, blends cached Values" --> CACHE
  NEW -. "new K,V appended to the cache" .-> CACHE

Why the cache holds Keys and Values

A Query is only used while its own token is the current position being processed. Once that step is done, no future token needs that old Query again. Keys and Values are different: every future token will compare its new Query against the old Keys and pull from the old Values. So the cache holds K and V, and discards each Query after it's used.

At generation step t, for each layer: compute the new token's q_t, k_t, v_t; append k_t, v_t to that layer's stored tensors; score q_t against all cached Keys; blend the cached Values.

scores_t = q_t [K₁ … K_t]^T / √d_k
output_t = softmax(scores_t) [V₁ … V_t]

import numpy as np

# grows by one entry per processed token
K_cache, V_cache = [], []

def generate_step(x_new, Wq, Wk, Wv):
    q = x_new @ Wq
    K_cache.append(x_new @ Wk)
    V_cache.append(x_new @ Wv)

    # everything cached so far
    K, V = np.stack(K_cache), np.stack(V_cache)
    scores = (q @ K.T) / np.sqrt(q.shape[-1])
    weights = np.exp(scores) / np.exp(scores).sum()

    return weights @ V   # blend of all cached Values

There's a separate cache at every layer. A token's Key and Value in layer 1 are not its Key and Value in layer 12: each layer sees a different hidden representation and uses different projection matrices (Part 11). The cache therefore stores K and V for every past position, in every layer, for every head.

What it saves, and what it costs

Benefit	Cost
Skips recomputing past tokens' K/V at each step. The new Query is still scored against every cached Key, which grows linearly with length.	Memory grows with context length: every layer must retain all past Keys and Values.
Lets each new token attend to the full history.	Reading a large cache eats memory bandwidth and eventually becomes the bottleneck.

Grouped-query attention shrinks the cache by giving each group of Query heads one shared Key–Value head (Ainslie et al., 2023). Multi-query attention is the extreme case of a single shared Key–Value head (Shazeer, 2019).

Previous ← The Complete Transformer Block

Next How GPT Trains and Generates →