Mid-level interview questionAsked at Netflix, Mistral AI · interview write-up
What does the KV cache store, why is using it exact rather than an approximation, and how much does it save?
Model answer
It stores every processed token's keys and values, per layer and head, so each decode step computes only the new token's.
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.
It's exact because of the causal mask: earlier tokens never depend on later ones, and greedy output was identical with and without the cache. After a 960-token prompt on our test machine, a step took 9.6 ms with the cache against 228.8 ms without, about 24 times faster.
Practise the rest
11 Context & KV cache questions, 492 in all, each with a model answer. Free.
More from Netflix interview questions, Mistral AI 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.
- 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?
- Here's a small transformer with a few bugs. Fix them, then add a KV cache for decoding.