The KV cache

When the model generates token t+1 it computes attention over all the earlier tokens. Without a cache, K and V are recomputed for the whole sequence every time — but they are identical to last time, since the earlier tokens do not change.

The KV cache: save K and V per layer and head. At the next token only the new position's K and V are computed and appended.

The effect: every new token costs O(t) instead of O(t²). For a 500-token answer that is the difference between seconds and minutes.

The price is memory. The cache grows linearly with the context length and often becomes what limits how many simultaneous users a GPU can handle.

import numpy as np

def softmax(z):
    z = z - z.max(-1, keepdims=True); e = np.exp(z); return e / e.sum(-1, keepdims=True)

class CachedAttention:
    def __init__(self, d):
        self.d = d; self.K = None; self.V = None
    def step(self, q, k, v):
        """One new token: q, k, v have shape (d,). Returns the attention output."""
        self.K = k[None] if self.K is None else np.vstack([self.K, k])
        self.V = v[None] if self.V is None else np.vstack([self.V, v])
        score = (self.K @ q) / np.sqrt(self.d)       # (t,)
        return softmax(score) @ self.V                # (d,)

rng = np.random.default_rng(0); d = 8
c = CachedAttention(d)
for t in range(5):
    out = c.step(rng.normal(size=d), rng.normal(size=d), rng.normal(size=d))
print(c.K.shape, out.shape)     # (5, 8) (8,)

The memory formula:

bytes = 2 (K and V) × layers × kv_heads × d_head × sequence length × batch × bytes per number

A 7B model (32 layers, 32 heads, d_head 128, fp16) at 8 192 tokens: 2 · 32 · 32 · 128 · 8 192 · 2 B ≈ 4.3 GB — for one user. Hence GQA (fewer KV heads), a quantised cache and a sliding window.