8.3 — Grouped-query and multi-query attention (GQA / MQA)

1. Why this session exists

The KV-cache analysis of 6.14 ended with a numerical wall: for LLaMA-7B at a $128k$-token context, the KV cache is $67$ GB per sequence in fp16 — larger than the model parameters themselves. The cache is dominated by the count $2 \cdot L \cdot h \cdot d_h \cdot T$, and the only freely-tunable factor among those (without changing the model's hidden size $d = h d_h$, the depth $L$, or the served context $T$) is the number of distinct K and V projections per layer.

The structural observation behind grouped-query attention (GQA) and multi-query attention (MQA): there is no rule that says each query head must have its own $W_K$ and $W_V$. We can share K and V across groups of query heads, paying a small quality cost and getting a linear reduction in KV-cache size and in decode-step memory bandwidth.

This session derives the construction, counts the savings, and places GQA / MQA as the standard choice in 2024+ frontier dense models (Llama-3, Mistral, Qwen, Gemma). MLA, the more aggressive 2024 alternative used by DeepSeek, is the subject of §8.5.

2. The design space: from MHA to MQA

Standard multi-head attention (MHA, 6.4) has $h$ query heads, $h$ key heads, and $h$ value heads — one of each per head index $\ell \in \{1, \dots, h\}$. Multi-query attention (MQA, Shazeer 2019[2], arXiv:1911.02150) shares one K head and one V head across all $h$ query heads. Grouped-query attention (GQA, Ainslie et al. 2023, arXiv:2305.13245) generalises: partition the $h$ query heads into $g$ groups of $h/g$ heads each, and use one K and one V projection per group, shared across the queries in that group.

Parameterise by the number of KV groups $g \in \{1, 2, \dots, h\}$:

| Variant | $g$ | K, V heads | Comment | |---|---:|---:|---| | MQA | $1$ | $1$ each | Most aggressive sharing | | GQA-2 | $2$ | $2$ each | Two K, two V across all query heads | | GQA-$g$ | $g$ | $g$ each | $g$ K-heads, $g$ V-heads | | MHA | $h$ | $h$ each | No sharing — recover standard MHA |

GQA-$h$ is MHA; GQA-$1$ is MQA. So the choice is a single integer $g$.

2.1 The forward pass, formally

Fix hidden size $d$, query heads $h$, group count $g \mid h$, per-head dimension $d_h = d/h$. Define the projection matrices:

The $W_Q$ matrix has the same shape as in MHA; the $W_K$ and $W_V$ matrices have width $g d_h$ instead of $h d_h = d$, i.e. they output $g$ (not $h$) K-heads. Given $X \in \mathbb{R}^{T \times d}$:

Reshape: $Q \to T \times h \times d_h$, $K \to T \times g \times d_h$, $V \to T \times g \times d_h$. Now per query head $\ell \in \{1, \dots, h\}$ we need to pick which K, V to use. Group head $\ell$ to KV group

i.e. the first $h/g$ query heads use K, V group $1$; the next $h/g$ query heads use group $2$; etc. Per query head, apply scaled dot-product attention (6.2) with the shared K and V of its group:

Concatenate the $h$ heads, project: $\text{GQA}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W_O \in \mathbb{R}^{T \times d}$. The output shape matches MHA. So GQA is a drop-in replacement for MHA at the level of input/output signatures; the change is interior to MHA.

3. KV-cache shrink: the headline win

The KV-cache equation of 6.14 §4 was

where $\text{bpe}$ is bytes per element. With GQA-$g$, the K, V tensors per layer have only $g$ (not $h$) heads. So

The cache shrinks by factor $h/g$. For MQA ($g = 1$), this is $h\times$ — a $32\times$ reduction for $h = 32$. For GQA-$8$ in a $h = 64$ model, this is $8\times$.

Numerical example. LLaMA-3-70B uses GQA-$8$ with $h = 64$, $L = 80$, $d = 8192$, $d_h = 128$. At $T = 32k$:

The MHA version would not fit on an 80 GB H100 even for batch size 1; the GQA-$8$ version leaves $\sim 70$ GB for model weights and activations. This is the operational reason for GQA's adoption: it makes long-context serving with batch parallelism feasible on standard 80 GB accelerators.

3.1 Compute count and parameter count

Parameter count of the attention sub-layer, with $W_Q \in \mathbb{R}^{d \times d}$, $W_K, W_V \in \mathbb{R}^{d \times g d_h}$, $W_O \in \mathbb{R}^{d \times d}$:

For $g = h$ (MHA) this is $4 d^2$ — the value of 6.4 §3. For $g = 1$ (MQA) it is $2 d^2 + 2 d^2 / h \approx 2 d^2$ — about half of MHA. Per-layer parameter savings are modest; the dominant win is the KV cache.

FLOPs per training-time forward pass per layer (compare 6.12):

• Projections ($Q, K, V,$ output): $2 d^2 T + 2 \cdot 2 d^2 (g/h) T + 2 d^2 T \approx 4 d^2 T (1 + g/(2h))$. For MQA this is $\approx 4 d^2 T$ instead of $8 d^2 T$ — half.
• Score matmul $Q K^\top / \sqrt{d_h}$ per head: $2 T^2 d_h$ per head, $h$ heads $= 2 T^2 d$. Independent of $g$.
• Mix with $V$: $2 T^2 d$. Independent of $g$.

So the dominant FLOPs (the $T^2 d$ terms) are the same for MQA, GQA, and MHA at training time. Decode-time FLOPs are dominated by the projections, which do shrink slightly with $g$. But the headline number is the KV-cache memory and the bandwidth to read it during decode.

3.2 Decode bandwidth, the operational metric

At decode, per generated token, we read $K \in \mathbb{R}^{T \times g d_h}$ and $V \in \mathbb{R}^{T \times g d_h}$ from HBM, perform a single-query attention of shape $1 \times T$, and discard most of the multiplied-out result. The bandwidth cost per layer per step is

For MQA, $g/h = 1/h$, so decode-step KV bandwidth drops by $h\times$. Combined with the (model-parameter-fixed) bandwidth to read the layer's weights, this gives a clean wall-clock decode speedup that scales with $h/g$ until the parameter-reading bandwidth dominates. For LLaMA-3-70B ($h = 64, g = 8$), the empirical decode-step speedup vs the hypothetical MHA-70B is $\sim 5$–$8\times$ depending on context length and batch.

4. The quality cost

The harder question: how much quality do we lose by tying K, V across query heads?

Ainslie et al. (2023)[1] measured this empirically by uptraining MHA checkpoints into GQA and MQA. The recipe:

1. Take an existing MHA-trained checkpoint with $h$ K and V heads.
2. Mean-pool the $h$ K-heads down to $g$ K-heads (and same for V): each group's K is the average of the $h/g$ original K's it absorbs.
3. Continue training for a small fraction ($\sim 5\%$) of the original training budget with the GQA mask in place.

Result: with $g = 8$, post-uptraining quality is within $\sim 0.1$ perplexity of the original MHA checkpoint on standard benchmarks. With $g = 1$ (MQA), quality drops a measurable $0.3$–$0.7$ perplexity. The interpretation: collapsing all heads to a single K, V pair loses some of the per-head specialisation observed in trained MHA models (6.4 §4), but $g = 8$ still has enough specialisation to recover most of the quality.

This is the empirical sweet spot the field converged on:

| Model | $h$ | $g$ | Variant | |---|---:|---:|---| | LLaMA-1, LLaMA-2-7B | 32 | 32 | MHA | | LLaMA-2-34B / 70B | 56 / 64 | 8 | GQA-$8$ | | LLaMA-3 (all sizes ≥ 8B) | 32 / 64 | 8 | GQA-$8$ | | Mistral-7B | 32 | 8 | GQA-$8$ | | Qwen-2 | 32 / 64 | 8 | GQA-$8$ | | Gemma-2 | 32 / 48 | 8 / 16 | GQA-$8$ / GQA-$16$ | | Falcon-180B | 64 | 1 | MQA | | PaLM (early) | — | 1 | MQA |

GQA-$8$ is the de facto 2024–2026 default for dense decoder-only models. MQA still appears in a few large dense models (Falcon-180B, PaLM); the empirical view is that at very large $h$, even $g = 1$ has enough effective capacity, but the $g = 8$ point gives a quality-cost insurance policy with little additional cache.

5. Uptraining vs from-scratch

A model can be either trained from scratch with GQA or uptrained from an MHA checkpoint. Three observations.

(a) From-scratch training is no harder than MHA training. Same losses, same recipes, same hyperparameters. The only change is the per-layer projection shapes. LLaMA-3 was trained from scratch with GQA-$8$; no pre-existing MHA checkpoint was involved.

(b) Uptraining from MHA works and is cheap. The mean-pool initialisation closes most of the gap with $\sim 5\%$ of original training compute. The technique is useful when an organisation already has an MHA model and wants to shrink its KV cache without retraining from scratch — LLaMA-2-70B's GQA variant was produced this way.

(c) The MHA → GQA transition is unidirectional. Once trained with $g$ groups, the model has internalised the group structure; you cannot recover the lost K, V diversity by "unpooling" the projections. This matters for serving: a model trained with $g = 8$ cannot be served with $g = 64$ to recover MHA quality at inference time.

6. Where GQA fits with the rest of Chapter 8

GQA is one of two attention-side KV-cache shrinkers; MLA (§8.5) is the other. They differ in their structural device:

• GQA shrinks the number of K, V heads — fewer projections, but each is full per-head-dimension.
• MLA shrinks the dimension of K, V via a low-rank latent — the number of heads is preserved (in the $Q$-side count), but K and V are stored compressed and decompressed at attention time.

They are not mutually exclusive (one could imagine GQA combined with MLA), but in practice no model uses both: GQA is the simpler, more widely-adopted option in dense Western LLMs, and MLA is DeepSeek's choice for their V2/V3/R1 lineage. We will derive MLA carefully in §8.5 to compare.

The other Chapter 8 sessions are orthogonal to the GQA/MLA choice: SwiGLU (§8.4) touches the FFN, MoE (§8.6) replaces the FFN with a router-plus-experts, Mamba (§8.7) replaces attention entirely with an SSM, hybrids (§8.8) mix Mamba and attention layers. Frontier dense models in 2026 typically combine GQA-$8$ + SwiGLU + RMSNorm + RoPE inside the block diagram of 8.1 §6.

7. What we keep going forward

Three takeaways for the rest of the chapter:

1. GQA-$g$ is parameterised by one integer. Replace $h$ K-heads and $h$ V-heads with $g$ K-heads and $g$ V-heads; tie groups of $h/g$ query heads to each KV group. Default $g = 8$.

2. The shrink is exactly $h/g$ in KV cache and decode-step KV bandwidth. That is the operational reason it became the 2024+ default.

3. Quality cost is small ($\lesssim 0.1$ perplexity) at $g = 8$ for dense models in the $7$B–$70$B range, after either from-scratch training or $\sim 5\%$-budget uptraining from MHA.

GQA is the first of three KV-cache treatments in Chapter 8. The next, SwiGLU (§8.4), shifts to the FFN side; MLA (§8.5) returns to attention with a more aggressive compression.