Mid-level interview questionAsked at OpenAI, Mercor · interview write-up
Here's a small transformer with a few bugs. Fix them, then add a KV cache for decoding.
Model answer
Usual planted bugs: the mask applied after softmax, scaling by d instead of head size, a missing transpose before merging heads. For the cache, each decode step computes q, k and v for the new token only, appends k and v, and attends over everything cached.
def step(x, w_qkv, w_o, h, cache): # x: (B, 1, d)
B, _, d = x.shape
q, k, v = (x @ w_qkv).split(d, -1)
sh = lambda t: t.view(B, 1, h, -1).transpose(1, 2)
q, k, v = sh(q), sh(k), sh(v) # (B, h, 1, d/h)
if "k" in cache:
k = torch.cat([cache["k"], k], dim=2)
v = torch.cat([cache["v"], v], dim=2)
cache["k"], cache["v"] = k, v
s = q @ k.transpose(-2, -1) / (d // h) ** 0.5
y = s.softmax(-1) @ v # no mask: all past
return y.transpose(1, 2).reshape(B, 1, d) @ w_o
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.
Interviewers check that cached decoding matches the full forward pass and that positions continue from the cache length.
Practise the rest
11 Context & KV cache questions, 492 in all, each with a model answer. Free.
More from OpenAI interview questions, Mercor 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?