Llama 3, Mistral, and Qwen3 have all adopted Grouped-Query Attention (GQA) as their default architecture. As a result, num_kv_heads in a model spec is no longer just an architectural parameter — it's an infrastructure variable that determines GPU memory budget and the upper bound on batch size. Choosing a model without considering the attention head count alongside your serving environment constraints will lead to unexpected bottlenecks in production.
How the Three Architectures Stack KV Heads
Multi-Head Attention (MHA) maintains H query, key, and value heads each. Multi-Query Attention (MQA), proposed by Shazeer (2019), reduces keys and values to a single head while keeping H query heads. The goal was to reduce HBM bandwidth consumption from reading the KV cache during the decode phase. Grouped-Query Attention (GQA) by Ainslie et al. (2023) sits in between — it groups key and value heads into G groups, so queries within each group share a single K·V pair. G=1 is equivalent to MQA; G=H is equivalent to MHA.
In serving, this difference has one concrete consequence: KV cache size scales linearly with num_kv_heads.
KV Cache Memory: Computing It Directly
KV cache bytes = 2 × num_kv_heads × head_dim × seq_len × num_layers × bytes_per_element
The factor of 2 accounts for storing both K and V. When all other parameters are equal, num_kv_heads is the only variable that matters.
Let's plug in Llama 3 8B: 32 query heads, 8 KV heads (GQA-8), head_dim 128, 32 layers, bf16 (2 bytes). We compare how the KV cache changes if the same model used MHA (num_kv_heads=32) or MQA (num_kv_heads=1).
For a single request under GQA-8:
2 × 8 × 128 × seq_len × 32 × 2 = 131,072 × seq_len bytes
| seq_len | MHA (32 KV heads) | GQA-8 (8 KV heads) | MQA (1 KV head) |
|---|---|---|---|
| 2,048 | 1.0 GB | 0.25 GB | 0.031 GB |
| 8,192 | 4.0 GB | 1.0 GB | 0.125 GB |
| 32,768 | 16.0 GB | 4.0 GB | 0.5 GB |
| 131,072 | 64.0 GB | 16.0 GB | 2.0 GB |
At seq_len 131,072, a single MHA request's KV cache fills an entire A100 80GB. GQA-8 uses 16 GB under the same conditions. But the per-request number matters less than what happens when you multiply it by batch size.
How KV Cache Savings Raise the Batch Size Ceiling
Llama 3 8B in bf16 occupies roughly 16 GB of weights. On an A100 80GB, after accounting for weights, activations, and system overhead, you have approximately 60 GB available for the KV cache. At seq_len 8,192:
| Architecture | KV cache per request | Max concurrent requests (60 GB budget) |
|---|---|---|
| MHA (32 KV heads) | 4.0 GB | ~15 |
| GQA-8 (8 KV heads) | 1.0 GB | ~60 |
| MQA (1 KV head) | 0.125 GB | ~480 |
GQA-8 can handle 4× more concurrent requests than MHA. The path from this batch size increase to decode throughput improvement is direct.
Decode is a memory-bound operation — every token generated requires reading the full K·V cache of all previous tokens from HBM. With the same HBM bandwidth, a larger batch increases arithmetic intensity per request, leaving GPU cores idle less of the time. KV cache savings → larger batches → higher arithmetic intensity → better tokens/s: these are all links in the same chain.
Measured inference time on T5-XXL from Ainslie et al. (2023): MHA 1.51s, MQA 0.24s (~6.3× faster), GQA-8 0.28s (~5.4× faster). The gap between GQA-8 and MQA is only 0.04s — essentially negligible.
Impact on TTFT and TPOT
GQA's benefits look different in prefill versus decode.
During prefill, GQA reduces attention FLOPs. Cutting num_kv_heads from 32 to 8 directly reduces the FLOPs for K·V projection and attention score computation. However, prefill is a compute-bound phase where FFN operations dominate the total time by a wide margin compared to attention. This is why measured TTFT improvements don't scale linearly with KV cache reduction ratio — switching to GQA won't give you a 4× reduction in TTFT.
Decode is a different story. Every generated token requires reading the K·V of all previous tokens from HBM. Fewer KV heads directly reduces the amount of data loaded, alleviating the HBM bandwidth bottleneck. The A100's HBM bandwidth is 2 TB/s, and with large batches and long sequences, KV cache loading becomes the bottleneck. TPOT (Time Per Output Token) directly measures decode speed, so GQA's improvements show up clearly here — and the effect grows with batch size.
vLLM selects the FlashAttention-2 GQA-optimized code path when num_kv_heads < num_query_heads. FA2 and later support grouped K·V processing natively, computing attention in group-sized chunks rather than repeating K·V tensors to match the query head count. Depending on the hardware, FA2 (default), FA3 (Hopper H100/H200), or FA4 (Blackwell B200) is selected.
Trade-offs in Choosing the Number of Groups
Average quality scores on T5-XXL from Ainslie et al. (2023): MHA 47.2, GQA-8 47.1, MQA 46.6. GQA-8 trails MHA by just 0.1 points, while MQA is 0.6 points lower. These gaps look small in aggregate, but the picture changes for long-context tasks.
Reducing K·V heads to 1 with MQA significantly limits each layer's ability to capture diverse attention patterns. This doesn't matter much for short sequences, but over contexts of tens of thousands of tokens, the reduced K·V expressiveness compounds and degrades retrieval and summarization quality. This is why Llama 3, despite supporting 128k context, chose GQA-8 over MQA — it cuts KV cache to one-quarter of MHA while holding quality loss to 0.1 points.
There's one more constraint that's easy to miss when choosing the number of GQA groups: the number of kv_heads must be compatible with your Tensor Parallel configuration. Specifically, kv_heads must be divisible by the number of TP GPUs. A model with kv_heads=8 on TP=8 distributes evenly to one head per GPU — that works, as does TP=4 or TP=2. But a model with kv_heads=4 on TP=8 can't be evenly divided, and you'll get a runtime error. Infrastructure teams frequently run into this constraint late, only after scaling up GPU count or adjusting TP settings. Verify kv_heads against your TP configuration at model selection time, not after deployment.