Grouped Query Attention
Grouped-Query Attention (GQA) is the attention architecture variant that shares Key and Value heads among groups of Query heads — reducing the KV cache memory footprint and inference cost by a factor equal to the group size, while retaining most of the quality of standard Multi-Head Attention (MHA), making it the dominant attention design in modern large language models including LLaMA 2/3, Mistral, and Gemma.
Attention Head Variants
| Variant | Query Heads | KV Heads | KV Cache Size | Quality |
|---|---|---|---|---|
| MHA (Multi-Head) | H | H | H × d_k × 2 | Best |
| GQA (Grouped-Query) | H | H/G (G groups) | H/G × d_k × 2 | Near-MHA |
| MQA (Multi-Query) | H | 1 | 1 × d_k × 2 | Slightly lower |
- MHA (original transformer): 32 query heads, 32 KV heads → full quality, full memory.
- MQA (Shazeer, 2019): 32 query heads, 1 KV head → 32x less KV cache, slight quality drop.
- GQA (Ainslie et al., 2023): 32 query heads, 8 KV groups → 4x less KV cache, negligible quality drop.
How GQA Works
Standard MHA (H=32 heads):
Q: 32 heads × d_k K: 32 heads × d_k V: 32 heads × d_k
Head i attends using Q_i, K_i, V_i
GQA (H=32 query, G=8 KV groups):
Q: 32 heads × d_k K: 8 groups × d_k V: 8 groups × d_k
Query heads 0-3 share KV group 0
Query heads 4-7 share KV group 1
...up to query heads 28-31 share KV group 7Memory and Compute Savings
- LLaMA-2 70B: 64 query heads, 8 KV heads (GQA with G=8).
- KV cache reduction: 8x compared to MHA → critical for long-context inference.
- For 4096-token context: KV cache drops from ~80 GB to ~10 GB for 70B model.
- Compute: KV projection compute reduced 8x (minor, since QKV projection is small relative to attention).
Why GQA Over MQA
- MQA (1 KV head) shows noticeable quality degradation on complex reasoning tasks.
- GQA (8 KV groups) matches MHA quality within noise on most benchmarks.
- GQA is a smooth interpolation: G=1 → MQA, G=H → MHA.
- Sweet spot: 4-8 KV groups for models with 32-128 query heads.
Models Using GQA
| Model | Query Heads | KV Heads | Ratio |
|---|---|---|---|
| LLaMA-2 70B | 64 | 8 | 8:1 |
| LLaMA-3 | 32 | 8 | 4:1 |
| Mistral 7B | 32 | 8 | 4:1 |
| Gemma | 16 | 1 (MQA) | 16:1 |
| Falcon 40B | 64 | 1 (MQA) | 64:1 |
| GPT-4 (rumored) | GQA variant | — | — |
Training Considerations
- GQA can be applied to existing MHA checkpoints via "uptraining" — merge KV heads by averaging, then fine-tune.
- Training from scratch with GQA: No special process — just configure fewer KV heads in architecture.
Grouped-Query Attention is the standard attention design for modern LLMs — by offering the near-optimal quality/efficiency tradeoff for KV cache reduction, GQA enables the practical deployment of large models at long context lengths where full MHA would be prohibitively memory-intensive.