GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, et al.2023EMNLP 2023
Read it on arxiv.org(opens in a new tab)Why this one
If you serve decoder models, the KV cache is not an implementation detail, it is often the product constraint. Multi-head attention spends memory bandwidth storing and reading many key and value heads. Multi-query attention cuts that down hard, but can cost quality. The useful idea here is the middle setting: keep many query heads, share fewer key and value heads, and convert an existing checkpoint instead of training from zero. People often talk about GQA as a model architecture feature, but for engineers its real importance is migration economics: you can buy a large fraction of the serving win without throwing away the pretrained model. Read this after Efficiently Scaling Transformer Inference and Efficient Memory Management for Large Language Model Serving with PagedAttention, because those explain the bottleneck this paper learns to remove.
What to take away
- GQA reduces KV cache size by sharing key and value heads across groups of query heads.
- The paper gives an uptraining path from MHA checkpoints instead of requiring a fresh pretraining run.
- The quality and latency tradeoff is tunable by choosing the number of key value groups.
Reads with
- Fast Transformer Decoding: One Write-Head is All You Need
gqa is the quality preserving midpoint between mha and this one kv head extreme
- Efficiently Scaling Transformer Inference
it explains why kv cache bandwidth dominates decoder serving cost
- Efficient Memory Management for Large Language Model Serving with PagedAttention
it attacks the same kv cache constraint from the runtime allocator side