Mid-level interview questionAsked at Mistral AI · interview write-up
What are multi-head and grouped-query attention, and what does GQA do to the KV cache?
Model answer
In multi-head attention every head has its own queries, keys and values. Grouped-query attention lets a group of query heads share one key-value head.
Scroll sideways to see the whole diagram
In multi-head attention, each head projects its own queries, keys and values. Every key and value is cached for every token.
Text version
- In multi-head attention, each head projects its own queries, keys and values. Every key and value is cached for every token.
- Grouped-query attention keeps all the query heads but lets each group read one shared key-value head. Multi-query attention is the extreme: one K, V for all.
- The cache shrinks in proportion to the key-value heads, which is why most large open models use GQA.
Ainslie et al. (2023), GQA; Llama 2 70B's configuration: 64 query heads, 8 key-value heads (Touvron et al. 2023, Section 2.2).
The cache holds keys and values per key-value head, so fewer of them means a proportionally smaller KV cache: Llama 2 70B's 64 query heads share 8, an 8× cut, with little quality loss.
Practise the rest
11 Context & KV cache questions, 492 in all, each with a model answer. Free.
More from 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.
- 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?
- Here's a small transformer with a few bugs. Fix them, then add a KV cache for decoding.