Grouped-Query Attention (GQA) and Multi-Query Attention (MQA): Mathematical Foundations, KV-Cache Bandwidth Reduction, Uptraining Recipes, and Tensor Parallelism Implications

The KV-Cache Bandwidth Wall

Autoregressive decoder inference is bottlenecked by memory bandwidth, not compute. At each decoding step, the model must reload the entire key-value (KV) cache from high-bandwidth memory (HBM) into the compute units. For a model with $H$ attention heads, sequence length $n$, head dimension $d_k$, and batch size $b$, the KV cache occupies $2 \times b \times H \times n \times d_k$ elements. With multi-head attention (MHA), this scales linearly with $H$.

Shazeer (2019) showed that the ratio of memory access to arithmetic operations during incremental decoding is $\Theta(n/d + 1/b)$ for MHA. When $n \approx d$ or $b \approx 1$, this ratio approaches 1, making memory bandwidth the dominant constraint on modern accelerators (Shazeer, 2019). The KV cache — not the model weights — becomes the primary bandwidth consumer at long context lengths.

Multi-Query Attention (MQA)

MQA, introduced by Shazeer (2019), shares a single key head and a single value head across all $H$ query heads. The query projections $P_q \in \mathbb{R}^{H \times d \times d_k}$ remain per-head, but the key and value projections collapse to $P_k \in \mathbb{R}^{d \times d_k}$ and $P_v \in \mathbb{R}^{d \times d_v}$.

Mathematical Formulation

For batched multi-query attention:

$$\begin{aligned} Q &= \text{einsum}(\text{"bnd, hdk->bhnk"}, X, P_q) \\ K &= \text{einsum}(\text{"bmd, dk->bmk"}, M, P_k) \\ V &= \text{einsum}(\text{"bmd, dv->bmv"}, M, P_v) \\ \text{logits} &= \text{einsum}(\text{"bhnk, bmk->bhnm"}, Q, K) \\ \text{weights} &= \text{softmax}(\text{logits} + \text{mask}) \\ O &= \text{einsum}(\text{"bhnm, bmv->bhnv"}, \text{weights}, V) \\ Y &= \text{einsum}(\text{"bhnv, hdv->bnd"}, O, P_o) \end{aligned}$$

The critical difference from MHA: $K$ and $V$ lose their head dimension $h$. During incremental decoding, the KV cache shape changes from $[b, H, m, d_k]$ to $[b, m, d_k]$, reducing KV memory by a factor of $H$.

Bandwidth Analysis

Shazeer's performance analysis shows the memory-to-compute ratio for incremental MQA becomes $\Theta(1/d + n/(dH) + 1/b)$. The offensive $n/d$ term from MHA is reduced by a factor of $H$. For a 32-head model, this is a 32x reduction in KV cache bandwidth pressure.

Experimental results on WMT14 En-De translation showed MQA with 8 heads achieved 27.5 BLEU vs 27.7 for MHA, while decoder inference latency dropped from 46 µs/token to 3.8 µs/token — a 12x speedup (Shazeer, 2019, Table 2).

Quality Trade-off

MQA's aggressive compression can degrade quality, especially on tasks requiring fine-grained information routing. Shazeer found MQA incurred minor quality degradation (0.2 BLEU on WMT14, 0.3 PPL on Billion-Word LM). However, training MQA from scratch proved unstable — loss spikes and divergence occurred during fine-tuning on long-input tasks.

Grouped-Query Attention (GQA)

Ainslie et al. (2023) proposed GQA as an interpolation between MHA and MQA. Query heads are partitioned into $G$ groups ($1 < G < H$), with each group sharing one key head and one value head. GQA-$G$ denotes $G$ groups; GQA-1 = MQA, GQA-$H$ = MHA.

Mathematical Formulation

For $H$ query heads and $G$ groups, each group contains $H/G$ query heads. The key and value projection matrices become $P_k \in \mathbb{R}^{G \times d \times d_k}$ and $P_v \in \mathbb{R}^{G \times d \times d_v}$. During attention computation, each group's key and value heads are broadcast to the $H/G$ query heads in that group.

The KV cache shape becomes $[b, G, m, d_k]$, reducing memory by a factor of $H/G$ compared to MHA. For GQA-8 with $H=32$, this is a 4x reduction — less aggressive than MQA's 32x, but often sufficient to alleviate the bandwidth bottleneck.

Uptraining Recipe

Ainslie et al. demonstrated that existing MHA checkpoints can be converted to GQA/MQA with minimal additional compute:

1. Checkpoint Conversion: Mean-pool the key and value projection matrices within each group. For MQA, pool all $H$ heads into one. For GQA-$G$, pool the $H/G$ heads within each group. Mean pooling outperforms selecting a single head or random initialization (Ainslie et al., 2023, Figure 4).
2. Uptraining: Continue pre-training for $\alpha = 5\%$ of original training steps (e.g., 50k steps for a 1M-step run) on the same data and optimizer settings. GQA converges stably; MQA shows higher variance and may require multiple fine-tuning runs.

The T5-XXL experiments showed GQA-8 uptrained with 5% compute achieved 47.1 average Rouge-1 vs 47.2 for MHA-XXL, with inference time 0.28s vs 1.51s — 5.4x faster while retaining 99.8% quality (Ainslie et al., 2023, Table 1).

Number of Groups

Figure 6 of Ainslie et al. shows inference time as a function of $G$. Going from $G=1$ (MQA) to $G=8$ adds modest overhead; $G=16$ approaches MHA latency. The authors selected $G=8$ as a favorable middle ground for T5-XXL (64 heads → 8 groups = 8 heads per group). For LLaMA-3-70B (64 heads, 8 KV heads), $G=8$ is the native configuration.

Tensor Parallelism Implications

Tensor parallelism (TP) shards attention heads across devices. In MHA with TP degree $T$, each device holds $H/T$ complete heads (Q, K, V, O projections). In GQA/MQA, the KV heads are fewer than query heads, creating an imbalance.

Megatron-LM Partitioning Strategy

Megatron-LM handles GQA by partitioning query heads across TP ranks while replicating KV heads (Megatron-LM, 2024). For GQA with $G$ KV heads and TP degree $T$:

• If $G \ge T$: Each rank gets $G/T$ KV heads and $H/T$ query heads. KV heads are sharded like MHA.
• If $G < T$: KV heads are replicated across ranks. Each rank holds all $G$ KV heads but only $H/T$ query heads. This increases memory per rank but avoids complex gather/scatter.

vLLM and SGLang implement similar logic: they replicate the smaller KV cache across TP workers when $G < T$, since the memory overhead is small compared to the communication cost of sharding.

Head Dimension Constraints

TP requires $H \bmod T = 0$ for even query head sharding. For GQA, $G \bmod T = 0$ is also desirable when $G \ge T$. LLaMA-3-70B (64 Q-heads, 8 KV-heads) works cleanly with TP=8 (8 Q-heads/rank, 1 KV-head/rank). TP=4 gives 16 Q-heads/rank, 2 KV-heads/rank. TP=16 would require KV replication since $8 < 16$.

Adoption in Modern LLMs

| Model | Architecture | Q-Heads | KV-Heads | Group Ratio | Source | |-------|--------------|---------|----------|-------------|--------| | PaLM / PaLM-2 | MQA | 48 | 1 | 48:1 | Chowdhery et al., 2022 | | LLaMA-1 7B/13B | MHA | 32/40 | 32/40 | 1:1 | Touvron et al., 2023 | | LLaMA-2 7B/13B | MHA | 32/40 | 32/40 | 1:1 | Touvron et al., 2023 | | LLaMA-2 70B | GQA | 64 | 8 | 8:1 | Touvron et al., 2023 | | LLaMA-3 8B/70B | GQA | 32/64 | 8/8 | 4:1 / 8:1 | Meta AI, 2024 | | Mistral 7B | GQA | 32 | 8 | 4:1 | Jiang et al., 2023 | | Mistral Large | GQA | 64 | 8 | 8:1 | Mistral AI, 2024 | | Qwen 2.5 7B/72B | GQA | 28/64 | 4/8 | 7:1 / 8:1 | Qwen Team, 2024 | | Gemma 2 9B/27B | GQA | 16/32 | 2/8 | 8:1 / 4:1 | Google, 2024 | | DeepSeek-V2 | MLA (latent) | 128 | 1 (compressed) | N/A | DeepSeek-AI, 2024 |

GQA has become the default for models >7B parameters. The 8:1 ratio (64:8 or 32:4) appears optimal for balancing KV cache reduction against representational capacity. MQA remains rare in open-weight models due to quality and stability concerns, though PaLM used it successfully at 540B scale.

Interaction with Other Inference Optimizations

PagedAttention

vLLM's PagedAttention manages KV cache in fixed-size blocks (typically 16 or 32 tokens) to eliminate fragmentation (Kwon et al., 2023). GQA/MQA reduce the per-block memory footprint proportionally to the group ratio. A GQA-8 block with 32 tokens at head dimension 128 uses 8 KB vs 64 KB for MHA (FP16), allowing 8x more blocks in the same GPU memory.

Speculative Decoding

Medusa, EAGLE, and other speculative decoding methods propose multiple tokens from a draft model, verified by the target model. The target model's KV cache bandwidth remains the bottleneck during verification. GQA/MQA directly accelerate this phase by reducing KV cache reads per verification step.

Prefix Caching

Automatic prefix caching (Anthropic, OpenAI, Google, vLLM APC) matches incoming prompts against cached KV prefixes. GQA/MQA reduce the memory cost of each cached prefix, enabling longer prefixes and higher cache hit rates at the same memory budget.

Training Stability Considerations

Ainslie et al. (2023) identified MQA training instability during fine-tuning on long-input tasks, characterized by loss spikes and immediate divergence. GQA-uptrained models were stable. The authors hypothesize that the single KV head in MQA creates gradient concentration that destabilizes optimization. GQA's multiple groups distribute gradients more evenly.

For practitioners converting MHA checkpoints:

• Use mean pooling for checkpoint conversion (not random init or head selection)
• Uptrain for 5% of original steps minimum; 10% shows diminishing returns
• Prefer GQA over MQA unless maximum speed is required and quality degradation is acceptable
• Monitor fine-tuning stability on long-context tasks; MQA may need multiple runs

Summary

| Aspect | MHA | GQA (e.g., 8 groups) | MQA | |--------|-----|----------------------|-----| | KV Heads | $H$ | $G$ | 1 | | KV Cache Reduction | 1x | $H/G$x | $H$x | | Quality | Baseline | ~MHA | Slightly degraded | | Inference Speed | Baseline | ~MQA | Fastest | | TP Friendliness | High | Medium (replicate if $G<T$) | High (replicate KV) | | Training Stability | Stable | Stable (uptrained) | Unstable (from scratch) | | Uptraining Compute | N/A | 5% | 5% |

GQA represents the practical sweet spot for production LLM serving: it delivers nearly all of MQA's bandwidth savings while retaining MHA's quality and training stability. The 5% uptraining recipe makes it accessible for teams with existing MHA checkpoints, and native GQA architectures (LLaMA-3, Mistral, Qwen, Gemma 2) now dominate the open-weight landscape.

Sources

• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150. https://arxiv.org/abs/1911.02150
• Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245. https://arxiv.org/abs/2305.13245
• Touvron, H. et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288. https://arxiv.org/abs/2307.09288
• Meta AI. (2024). Introducing Meta Llama 3. https://ai.meta.com/blog/meta-llama-3/
• Jiang, A. Q. et al. (2023). Mistral 7B. arXiv:2310.06825. https://arxiv.org/abs/2310.06825
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. arXiv:2309.06180. https://arxiv.org/abs/2309.06180
• DeepSeek-AI. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434. https://arxiv.org/abs/2405.04434
• Google. (2024). Gemma 2: Improving Open Models at 9B and 27B Parameters. arXiv:2408.00118. https://arxiv.org/abs/2408.00118
• Chowdhery, A. et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311. https://arxiv.org/abs/2204.02311