Senior interview questionAsked at Applied Intuition · interview write-up
Design and implement a KV cache for autoregressive decoding in a Transformer: tensor shapes for prefill and per-step decode, per-layer caches, batching, memory limits and beam search.
Model answer
class KVCache:
def __init__(self, L, B, H, T_max, dh, dtype, dev):
shape = (L, B, H, T_max, dh)
self.k = torch.zeros(shape, dtype=dtype, device=dev)
self.v = torch.zeros_like(self.k)
def update(self, layer, k, v, start):
# k, v: (B, H, n, dh); n = prompt len, then 1
end = start + k.size(2)
self.k[layer, :, :, start:end] = k
self.v[layer, :, :, start:end] = v
return (self.k[layer, :, :, :end],
self.v[layer, :, :, :end])
def reorder(self, beam_idx):
# beam search: rows follow surviving beams
self.k = self.k.index_select(1, beam_idx)
self.v = self.v.index_select(1, beam_idx)
Prefill writes the whole prompt in one pass; each decode step writes one position and attends over everything so far. Batch with left padding plus an attention mask so positions line up.
Scroll sideways to see the whole diagram
Prefill runs the prompt once and computes the keys and values of all 4 prompt tokens, in every layer and every head. They go into the cache.
Text version
- Prefill runs the prompt once and computes the keys and values of all 4 prompt tokens, in every layer and every head. They go into the cache.
- Each decode step computes the key and value of the new token only, and appends them to the cache.
- The new token's query still attends to every earlier key. Those keys and values are read back from the cache, not recomputed.
- Earlier tokens' keys and values never change, because of the causal mask, so the cache is exact. With a 4-token prompt and 3 new tokens, that is 7 key-value computations with the cache and 22 without.
Counts from the KV-cache lab in lesson 5: a 4-token prompt and 3 new tokens.
Memory per token is 2 × layers × KV heads × head dim × bytes, so cap batch times length. Beyond that: GQA, quantized caches, or paged blocks that avoid reserving T_max per sequence.
Practise the rest
11 Context & KV cache questions, 492 in all, each with a model answer. Free.
More from Applied Intuition interview questions.
More Context & KV cache interview questions
- What is a context window, and why can't GPT-2 small take a 1,025th token even with plenty of memory?
- Users say your summariser 'takes ages to start' on long documents but then writes quickly. Explain what's happening in terms of prefill and decode.
- What does the KV cache store, why is using it exact rather than an approximation, and how much does it save?
- You're sizing GPUs for a chat service. How do you estimate KV-cache memory, and which levers would you pull if it doesn't fit?
- Request A has a 960-token prompt and 10 output tokens. Request B has a 16-token prompt and 200 output tokens. Same model and machine: which finishes first, and what does that tell you about a chat app that resends the full history every turn?
- What are multi-head and grouped-query attention, and what does GQA do to the KV cache?