Sharing K/V heads to reduce cache size
Multi-Head Attention (MHA) gives each head its own key and value projections. With 8 heads, you compute 8 independent K-V caches during generation. Multi-Query Attention (MQA) uses a single shared K-V, reducing cache size by 8x. Grouped-Query Attention (GQA) is a middle ground: two heads share one K-V pair.
The KV cache size is proportional to sequence length and the number of K-V heads. Sharing K-V dramatically reduces memory overhead, enabling longer sequences or larger batch sizes on the same hardware.
Quality trade-off and recovery
MQA sacrifices some quality compared to MHA; the single shared K-V is a bottleneck. GQA reduces cache size less but usually matches MHA quality more closely. For smaller models, the quality gap is acceptable. For large models, GQA is often preferred.
Fine-tuning a pretrained MHA model to use MQA can recover some of the quality gap, but it requires retraining and is only partially effective.
When to use each variant
MHA is the standard for training and small models. MQA is useful for inference on hardware with limited memory (mobile, edge). GQA is increasingly popular for large models, offering a good balance: significant cache reduction with minimal quality loss. Models like Llama 2 use GQA for this reason.