← All papers

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón, Sanghai · 2023 · EMNLP 2023

EfficiencyRead on arXiv

Proposes Grouped Query Attention, an interpolation between multi-head and multi-query attention that groups query heads to share KV heads. Achieves quality close to MHA with speed close to MQA, and can be obtained by uptraining existing MHA checkpoints with only 5% of original compute.

Key Idea

Multi-Head Attention (MHA) gives each query head its own key and value head, maximum expressiveness but large KV cache. Multi-Query Attention (MQA) uses a single KV head shared across all query heads, giving fast inference but degraded quality. Grouped Query Attention (GQA) splits query heads into groups, with each group sharing one KV head. It's the sweet spot: near-MHA quality with near-MQA speed.

How It Works

  • Standard MHA: n_heads query heads, each with its own K and V head → KV cache = 2 × n_heads × d_head per token
  • MQA: n_heads query heads, all sharing 1 K and 1 V head → KV cache = 2 × d_head per token
  • GQA: n_heads query heads split into g groups, each group shares 1 KV head → KV cache = 2 × g × d_head per token
  • GQA with g=1 is MQA; GQA with g=n_heads is MHA
  • Uptraining: Convert an existing MHA checkpoint to GQA by mean-pooling KV heads within each group, then continue training for ~5% of original compute

Why It Matters

  • Adopted by nearly every major open LLM: Llama 2, Llama 3, Mistral, Gemma, Phi, GQA is the de facto standard
  • Practical KV cache reduction: 4-8× smaller cache than MHA with minimal quality loss
  • Uptraining from existing checkpoints: Don't need to pretrain from scratch, convert MHA models to GQA cheaply
  • Established the design point that all subsequent attention efficiency work (MLA, etc.) builds upon

Key Takeaways for Interviews

  • GQA = share KV heads across groups of query heads, interpolates between MHA (all separate) and MQA (all shared)
  • The number of KV groups g is a tunable knob: more groups = higher quality but larger cache
  • Typical configs: Llama 2 70B uses 8 KV heads with 64 query heads (8 groups of 8)
  • Uptraining trick: mean-pool existing KV heads within each group, then finetune, only 5% of pretraining compute
  • In system design: GQA directly determines KV cache memory budget, which determines max batch size and thus serving throughput