Tutorial 12: KV Cache and Autoregressive Generation¶
When generating text token by token, you could re-run the entire prompt through the model at every step. For a 1000-token prompt, step 50 would recompute all 1000 tokens' attention — wasteful, since the first 499 tokens' K and V have not changed. The KV cache stores them so we only process the new token at each step.
1. The Problem: Redundant Computation¶
Without caching, generating N tokens from a prompt of length P costs:
Step 1: process P+0 tokens → 1 new token
Step 2: process P+1 tokens → 1 new token
...
Step N: process P+N-1 tokens → 1 new token
Total attention computation: P + (P+1) + ... + (P+N-1) ≈ P×N + N²/2
For P=1000, N=500: ~500K + 125K = 625K attention operations.
With caching, we only process the new token at each step (and append its K/V to the cache):
Step 1: process 1 new token, attend to cached K/V → 1 new token
Step 2: process 1 new token, attend to cached K/V → 1 new token
...
Step N: process 1 new token, attend to cached K/V → 1 new token
Total: P (prefill) + N (decode) ≈ P + N
For P=1000, N=500: ~1500 operations. That is a 400× speedup.
2. How the KV Cache Works¶
During attention, each layer computes K and V for the current input:
K = linear(x, W_k) → shape: (batch, num_kv_heads, seq_len, head_dim)
V = linear(x, W_v) → shape: (batch, num_kv_heads, seq_len, head_dim)
The KV cache stores these and returns the full history (past + current):
# During prefill (processing the full prompt):
cache.update_and_fetch(K_new, V_new)
# Returns: K_full = [K_new], V_full = [V_new]
# The cache now holds all K/V for the prompt
# During decode (processing one new token):
K_new, V_new = attention(x) # shape: (batch, num_kv_heads, 1, head_dim)
cache.update_and_fetch(K_new, V_new)
# Returns: K_full = [K_past; K_new], V_full = [V_past; V_new]
# The cache now holds all K/V for the prompt + generated tokens
The attention computation then uses the full K and V to compute the output for the new token only.
3. The Dense Cache Implementation¶
The simplest cache concatenates new K/V with the existing history:
class TinyKvFullCache:
def __init__(self):
self.key_values = None # (keys, values) tuple, or None
self.offset = 0 # total tokens stored
def update_and_fetch(self, key, value, mask_length=None, mask=None):
"""
Append new K/V to the cache and return the full history.
Args:
key: New key tensor. Shape: (batch, num_heads, new_len, head_dim)
value: New value tensor. Same shape as key.
Returns:
(full_keys, full_values, seq_len, mask)
"""
if self.key_values is None:
# First call: just store the input
self.key_values = (key, value)
B, H, S, D = key.shape
self.offset = S
return key, value, self.offset, mask
# Subsequent calls: concatenate with existing history
prev_keys, prev_values = self.key_values
new_keys = mx.concat([prev_keys, key], axis=2) # append along seq dim
new_values = mx.concat([prev_values, value], axis=2)
self.key_values = (new_keys, new_values)
self.offset += key.shape[2]
return new_keys, new_values, self.offset, mask
Shape Tracking¶
Prefill: key shape (1, 8, 100, 128) → offset = 100
Decode 1: key shape (1, 8, 1, 128) → offset = 101
Decode 2: key shape (1, 8, 1, 128) → offset = 102
...
Each decode step adds one token to the sequence dimension.
4. The Rewrite Problem¶
Concatenation is simple but has a hidden cost: memory copying. Every time we append, we create a new tensor that copies all previous K/V data:
Step 1: allocate 100 × D bytes
Step 2: allocate 101 × D bytes (copy 100 + write 1)
Step 3: allocate 102 × D bytes (copy 101 + write 1)
...
After N steps, the total bytes copied are:
This is O(N²) in memory traffic. For long sequences, this dominates compute. This is the motivation for paged attention (Tutorial 13) — it avoids copying by writing new K/V into pre-allocated pages.
5. The Autoregressive Generation Loop¶
Here is the full generation loop using the KV cache:
def stream_generate(model, prompt_tokens, *, max_new_tokens=256,
sampler=None, eos_token_ids=(), cache=None):
"""
Generate tokens one at a time using autoregressive decoding.
Args:
model: The language model (must support __call__ and create_kv_cache).
prompt_tokens: List of input token IDs.
max_new_tokens: Maximum number of tokens to generate.
sampler: Function that maps logprobs to token IDs.
eos_token_ids: Set of token IDs that signal generation should stop.
cache: Pre-initialized KV cache (created if None).
Yields:
One token ID per iteration.
"""
tokens = [int(t) for t in prompt_tokens]
# Create the cache (one per layer)
if cache is None:
cache = model.create_kv_cache("dense")
eos = {int(e) for e in eos_token_ids}
offset = 0
for step in range(max_new_tokens):
if step == 0:
# Prefill: process the entire prompt at once
chunk, chunk_len = tokens, len(tokens)
else:
# Decode: process only the last (new) token
chunk, chunk_len = [tokens[-1]], 1
input_ids = mx.array([chunk])
# Forward pass through the model
# logits_to_keep=1: we only need the last position's prediction
logits = model(input_ids, offset=offset, cache=cache, logits_to_keep=1)
# Convert logits to probabilities (log-space for numerical stability)
logprobs = logits[:, -1, :] - mx.logsumexp(
logits[:, -1, :], axis=-1, keepdims=True
)
# Sample the next token
if sampler is not None:
next_id = int(sampler(logprobs)[0])
else:
next_id = int(mx.argmax(logprobs, axis=-1)[0])
yield next_id
# Stop if we hit an end-of-sequence token
if next_id in eos:
return
tokens.append(next_id)
offset += chunk_len
Key Details¶
logits_to_keep=1: During decode, we only have one new token, so we only
need one logit vector. This saves memory and compute.
offset: Tells the model (and RoPE) what absolute position this token is at.
For the prefill, offset=0. For decode step k, offset=prompt_len+k.
The first step is special: During prefill, we process all prompt tokens in parallel. During decode, we process one token at a time. The KV cache handles this transparently — it does not care whether the input has 1000 tokens or 1.
6. Prefill vs Decode¶
| Aspect | Prefill | Decode |
|---|---|---|
| Input length | Full prompt (e.g., 1000 tokens) | 1 token |
| Compute bound | Yes (large matmuls) | No (small matmuls) |
| Memory bound | No | Yes (loading weights for 1 token) |
| Mask | Causal (lower-triangular) | None (single token attends to all) |
| KV cache | Populated | Read + append |
| Batch efficiency | High (all tokens parallel) | Low (1 token per step) |
The prefill is compute-bound — we are doing large matrix multiplications for all prompt tokens simultaneously. The decode is memory-bound — we are loading the entire model's weights to process a single token. This asymmetry is why serving LLMs efficiently is challenging.
7. Continuous Batching with the KV Cache¶
When serving multiple users, each request has its own set of KV caches (one per layer). The scheduler manages these caches:
Request A: caches[0]=cache_A_0, caches[1]=cache_A_1, ...
Request B: caches[0]=cache_B_0, caches[1]=cache_B_1, ...
During each decode step, the scheduler batches together the new tokens from all active requests and runs them through the model in one forward pass:
# Decode step for 3 active requests
tokens = mx.array([
[token_A], # request A's next token
[0], # slot unused (padding)
[token_B], # request B's next token
])
offsets = [offset_A, 0, offset_B]
# One forward pass processes all active requests
logits = model(tokens, offset=offsets, cache=batched_cache)
The BatchingKvCache handles the complexity of assembling per-request caches
into batched tensors. This is what makes continuous batching possible — different
requests at different sequence lengths can share a single batch.
8. Summary¶
- The KV cache stores past K and V vectors, avoiding recomputation.
- Dense cache: simple concatenation, but O(N²) memory traffic.
- Prefill processes the full prompt; decode processes one token.
- The offset mechanism tells RoPE the absolute position of each token.
- Continuous batching manages per-request caches for multiple users.