Grouped Query Attention (GQA)

GQA occupies a precise middle ground between Multi-Head Attention (MHA) and Multi-Query Attention (MQA). In MHA, the number of key-value heads equals the number of query heads (H=K=V). In MQA, all query heads share a single key-value head (K=V=1). GQA generalizes this by dividing query heads into G groups, where each group shares one key-value head. This creates a tunable knob: setting G=1 yields MQA, while G=H yields MHA. Typical configurations use G=4 or G=8, reducing the KV cache size by a factor of H/G compared to MHA while preserving most of the quality.

The primary motivation for GQA is reducing the memory footprint of the KV cache during autoregressive decoding. For a model with H query heads, the KV cache stores 2 tensors (keys and values) per head per token. GQA with G groups reduces this to 2G tensors per token. This yields a memory reduction factor of H/G. For a 70B parameter model with 64 heads and G=8, the KV cache shrinks by 8x. This directly translates to larger batch sizes, longer context windows, and lower per-user serving costs without requiring architectural retraining from scratch.

GQA models are not trained from scratch. The standard procedure is uptraining: starting from a pre-trained MHA checkpoint and converting it to GQA. This is done by mean-pooling the key and value projection weights of the heads within each group. The model then undergoes a small fraction (typically 5-10%) of the original training steps to adapt. This uptraining is remarkably sample-efficient; the model quickly recovers quality because the underlying representations are already well-formed. This allows teams to retrofit existing large models with GQA for inference savings without the cost of full pre-training.

Empirical evaluations demonstrate that GQA achieves near-identical performance to MHA across standard benchmarks while approaching the speed of MQA. Key findings include:

• Perplexity: GQA with G=8 matches MHA perplexity on language modeling tasks.
• Downstream tasks: Negligible degradation on MMLU, HellaSwag, and reading comprehension.
• Long-context: GQA maintains quality better than MQA on sequences exceeding 10k tokens. The slight grouping constraint acts as a mild regularizer, and the shared key-value heads learn more robust representations that generalize well across the query heads within their group.

GQA has become the de facto standard for large-scale production language models due to its inference efficiency. Notable implementations include:

• Llama 2 70B: Uses GQA with 8 key-value heads (G=8) across 64 query heads.
• Llama 3: Continues the GQA architecture across all model sizes.
• Mistral: Employs GQA with sliding window attention for efficient long-context processing.
• Gemma: Google's open models use GQA to balance quality and serving costs. The widespread adoption validates GQA as the optimal trade-off point for models deployed at scale.

From a mechanistic interpretability perspective, GQA constrains the OV circuit formation. In MHA, each head can develop a specialized copy operation. With GQA, the shared value projection forces heads within a group to read from the same source representation. This encourages complementary query specialization: heads in a group learn distinct attention patterns (QK circuits) while operating on identical value information. Analysis shows this leads to more interpretable head clusters, where a group collectively represents a concept while individual heads route information to different token positions based on context.