You migrate an MHA decoder (32 query heads, 32 KV heads, head_dim 128) to grouped query attention with num_kv_heads=8, keeping 32 query heads and head_dim unchanged. Predict the factor by which the KV cache shrinks, and how many query heads now share each KV head.
KV cache scales with kv_heads, not query heads; 32 KV heads to 8 is a 4x cache shrink, with each KV head shared by a group of 4 query heads.
Picture 32 chefs in a kitchen, each holding a private copy of the same recipe book. The recipe books take up most of the counter space. Now you say: every group of 4 chefs shares one recipe book. You still have 32 chefs cooking 32 different dishes, but you only need 8 recipe books on the counter. The cooking happens at the same speed because the chefs are still independent, but you have freed up most of the table. Grouped-query attention does the same trick: the query heads stay, but they share keys and values in groups. The cache that holds keys and values shrinks by the same factor as the number of shared groups, which is why an 8-KV-head model uses a quarter of the cache memory of a 32-KV-head model.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Grouped-query attention is the single most important inference-side architecture change of the 2023 to 2026 era. It is also one of the easiest to compute about: the per-token KV cache size scales linearly with the number of KV heads, and going from full multi-head attention (32 KV heads) to GQA with 8 KV heads is a clean 4x cache reduction at constant head_dim.
The question asks for two numbers. The shrink factor: 4. The group size: 4. Both come from the same ratio, 32 / 8, but they answer different questions. The shrink factor compares cache before and after. The group size says how many query heads now share a single K/V head when computing attention.
This walkthrough derives both from the cache formula, explains why the quality story is good, and places GQA on the spectrum from full MHA to MQA to MLA in modern serving stacks.
The KV cache formula and why query heads vanish
The canonical per-token per-layer cache formula is:
The leading 2 covers K and V together. H_kv is the number of KV heads. d_head is the per-head dimension. b is bytes per value (2 for fp16/bf16, 1 for fp8, 0.5 for int4).
Queries do not appear. The reason is mechanical: at each decoding step, the new token computes its own Q from a single token's worth of activations, attends over all cached K and V, and writes its Q nowhere. Q is consumed and discarded. K and V from past tokens are read again on every subsequent step, so they have to live in cache. With n_query_heads independent of n_kv_heads (which is exactly what GQA enables), the cache cost decouples from query count.
Plugging in the question's numbers: at 32 KV heads, head_dim 128, fp16: 2 x 32 x 128 x 2 = 16384 bytes per token per layer. At 8 KV heads: 2 x 8 x 128 x 2 = 4096. The ratio is 4. That is the cache shrink factor.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Llama 3 8B and 70B both ship 8 KV heads with 32 and 64 query heads respectively (groups of 4 and 8).
- Mistral 7B uses GQA with 8 KV heads to fit 32k context on consumer GPUs.
What an interviewer would ask next. Try answering before peeking at the approach.
QHow does GQA interact with tensor parallelism across GPUs?
Query heads are typically sharded across the TP group, one head per rank or a small group. KV heads are also sharded but at a coarser granularity; you want n_kv_heads divisible by the TP degree, or you replicate KV across some ranks. Choosing 8 KV heads at TP=8 lets each rank hold exactly one KV head with no replication.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Computing the reduction from query-head count rather than KV-head count. Query heads do not enter the cache; only K and V tensors do, so the cache shrink factor is the kv_heads ratio, not the query_heads ratio.
60 second bullets to scan on the way to the call.
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.