KV-Cache: Why Generation Doesn't Recompute Everything

During autoregressive generation, the model produces one token at a time. Naively, you'd run the whole sequence through every layer for each new token — O(n²) compute per token, prohibitively slow. The KV-cache avoids this. The idea is mechanical and beautiful: K and V vectors for past tokens never change once computed, so cache them.

How it works

At step t, you only need to compute Q for the new token. K and V for the new token are computed and appended to the cache. Attention for the new token is then a single matmul against the cached K and V — Q has shape (1, d_head), K_cache has shape (t, d_head), so the score is shape (1, t) and the output is (1, d_v). Per-token compute drops from O(n²) to O(n).

The cost: memory

The cache stores 2 × n_layers × n_kv_heads × seq_len × d_head × bytes_per_param bytes. For Llama 3.3 (70B, 80 layers, 8 KV heads, d_head=128) at 128K context in FP16: roughly 32 GB. This is huge — often comparable to or larger than the model weights. It's the main reason inference at long context is hard, and the main reason GQA exists.

Code

KV-cache pseudocode·python

class KVCache:
    def __init__(self, max_len, n_kv_heads, d_head, n_layers, dtype, device):
        # one cache per layer; here just one for clarity
        self.K = torch.zeros(max_len, n_kv_heads, d_head, dtype=dtype, device=device)
        self.V = torch.zeros(max_len, n_kv_heads, d_head, dtype=dtype, device=device)
        self.length = 0
    def append(self, k_new, v_new):
        # k_new, v_new: (1, n_kv_heads, d_head)
        self.K[self.length] = k_new
        self.V[self.length] = v_new
        self.length += 1
    def get(self):
        return self.K[:self.length], self.V[:self.length]

def step_with_cache(token_id, cache, model):
    q = model.q_proj(model.embed(token_id))    # (1, n_q_heads, d_head)
    k = model.k_proj(model.embed(token_id))    # (1, n_kv_heads, d_head)
    v = model.v_proj(model.embed(token_id))
    cache.append(k, v)
    K, V = cache.get()
    # attention against full cached K, V — this is O(t), not O(t^2)
    scores = q @ K.transpose(-2, -1) / d_head**0.5
    weights = torch.softmax(scores, dim=-1)
    return weights @ V

KV cache size estimator·python

def kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, bytes_per_param=2):
    # 2 because we store both K and V
    return 2 * n_layers * n_kv_heads * d_head * seq_len * bytes_per_param

# Llama 3.3 70B at 128K context, FP16
b = kv_cache_bytes(80, 8, 128, 128_000, bytes_per_param=2)
print(f"{b / 1e9:.1f} GB")     # ~32 GB

KV-Cache: Why Generation Doesn't Recompute Everything

How it works

The cost: memory

Code

External links

Exercise

Progress

Comments 0