GQA: Grouped-Query Attention — Bridging Multi-Head Quality and Multi-Query Speed

Review date: 2026-05-27 Review author: Zhongzhu Zhou Paper reviewed: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints Paper authors: Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai arXiv: 2305.13245 Status/Venue: EMNLP 2023

Short Answer

Grouped-Query Attention (GQA) is a structural attention variant that sits on a continuum between the classical Multi-Head Attention (MHA) of the original Transformer and the aggressive Multi-Query Attention (MQA) of Shazeer 2019. Instead of either giving every query head its own key-value (KV) head (MHA, $H$ KV heads) or collapsing all queries onto a single shared KV head (MQA, $1$ KV head), GQA groups the $H$ query heads into $G$ groups, where each group shares one KV head. The cardinalities $G = H$ and $G = 1$ recover MHA and MQA exactly; intermediate $G$ values give intermediate cost-quality tradeoffs.

The paper makes three contributions:

1. It introduces the GQA architecture itself, which is now the de-facto default in production LLMs (LLaMA 2 70B, LLaMA 3, Mistral, Falcon, Qwen, DeepSeek-V2 base) precisely because it matches the quality of MHA while delivering nearly the decode-time bandwidth savings of MQA.
2. It proposes an uptraining procedure that converts an existing MHA checkpoint into a GQA (or MQA) checkpoint at only 5% of the original pretraining compute, using a principled mean-pooling of the per-head KV projections within each group as the initialization.
3. It runs the experiments that establish the canonical tradeoff curve: with $H=64$ query heads and $G=8$ groups, GQA-8 on T5 XXL is 5.4× faster than MHA-XXL at inference and only 0.1 average points worse on a suite of summarization, translation and QA benchmarks.

The single most important quantitative result of the paper is the T5 XXL row: MHA = 47.2 average score / 1.51 s per sample; MQA = 46.6 / 0.24 s; GQA-8 = 47.1 / 0.28 s. The 0.04 s gap to MQA buys 0.5 average points of quality back — a steep and well-positioned point on the Pareto frontier.

If you are deciding what attention variant to use for a serving-oriented model today, you should pick GQA with $G$ in $\{4, 8\}$ unless you have a specific reason not to. The remainder of this review explains, in painstaking detail, why this is true.

Prerequisites: What You Need to Know First

To understand GQA we have to first understand what MHA actually does at decode time, why the KV cache is the dominant memory object during autoregressive generation, and what “memory bandwidth bottleneck” means in a precise quantitative sense. This section builds that foundation. Readers who already have it can skip ahead to the What This Paper Does section.

1. The Transformer Decoder and Autoregressive Generation

A decoder-only Transformer (the architecture behind GPT, LLaMA, Mistral) generates text one token at a time. Given a prompt of length $P$, the model first runs a single prefill forward pass that computes all hidden states $h_1, \ldots, h_P$ in parallel. After that, generation proceeds in a decode loop:

$$y_{t+1} = \mathrm{argmax}_v \; \mathrm{softmax}\big(W_{out} h_t\big)_v$$

for $t = P, P+1, \ldots$ until an end-of-sequence token is sampled. At each decode step the model receives one token id, embeds it, runs it through every layer, projects to vocabulary, samples the next token, and appends.

This is profoundly different from training. In training, $L$ tokens are processed in one parallel forward pass. In decode, $L$ tokens require $L$ sequential forward passes, each of which has batch size 1 in the sequence dimension. The arithmetic intensity (FLOPs per byte loaded from memory) collapses, and the workload becomes memory bound instead of compute bound.

flowchart LR
    subgraph Prefill["Prefill: parallel over P tokens"]
        P1[Token 1] --> H1[h_1]
        P2[Token 2] --> H2[h_2]
        P3[Token ...] --> H3[h_...]
        PP[Token P] --> HP[h_P]
    end
    subgraph Decode["Decode: one token at a time"]
        D1[h_P] --> Y1[y_P+1]
        Y1 --> D2[Forward pass]
        D2 --> Y2[y_P+2]
        Y2 --> D3[Forward pass]
        D3 --> Y3[y_P+3]
    end
    HP -.last hidden state.-> D1

The decode loop is where production cost dominates. A user types 50 tokens of prompt and the model generates 500 tokens of response: the prefill is amortized across the 500 decode steps, and the decode steps individually each have to be fast.

2. Multi-Head Attention — The Standard Building Block

The Transformer’s attention block takes an input sequence $X \in \mathbb{R}^{L \times d_{model}}$ and produces an output of the same shape. With $H$ heads, each of dimension $d_k = d_{model}/H$, the standard MHA is:

$$\mathrm{MHA}(X) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_H) W^O$$ $$\mathrm{head}_h = \mathrm{softmax}\!\left(\frac{Q_h K_h^\top}{\sqrt{d_k}}\right) V_h$$ $$Q_h = X W_h^Q, \quad K_h = X W_h^K, \quad V_h = X W_h^V$$

with parameter matrices $W_h^Q, W_h^K, W_h^V \in \mathbb{R}^{d_{model} \times d_k}$ and output projection $W^O \in \mathbb{R}^{H d_k \times d_{model}}$.

Three properties matter for what follows:

• Per-head independence. Each head has its own $W^Q$, $W^K$, $W^V$. The output projection $W^O$ is the only place where heads interact.
• Symmetric parameter count. Q, K, V projections each have $H d_k d_{model}$ parameters per layer. They are symmetric: dropping K and V costs us 2/3 of the projection parameters but, more importantly for decode, 2/2 of the KV cache.
• Inner-product structure. Each $\mathrm{head}_h$ is a softmax over inner products. Heads can learn to attend to different positions, different syntactic relations, different topics. That diversity is the empirical reason MHA works better than single-head attention.

Below is the full per-step MHA decode algorithm, written explicitly so we can later compare its memory access pattern to MQA and GQA.

Algorithm 1: Multi-Head Attention (decode step t, one layer)
Input:
  x_t           in R^{1 x d_model}   -- current token's hidden state
  W_h^Q, W_h^K, W_h^V for h = 1..H
  K_cache       in R^{L x H x d_k}   -- previous keys
  V_cache       in R^{L x H x d_k}   -- previous values
  W^O           in R^{(H d_k) x d_model}

 1.  for h = 1..H do
 2.      q_h        = x_t W_h^Q                         # [1 x d_k]
 3.      k_h_new    = x_t W_h^K                         # [1 x d_k]
 4.      v_h_new    = x_t W_h^V                         # [1 x d_k]
 5.      K_cache[t, h, :] = k_h_new                     # append
 6.      V_cache[t, h, :] = v_h_new                     # append
 7.      scores_h   = q_h (K_cache[:t+1, h, :])^T / sqrt(d_k)   # [1 x (t+1)]
 8.      attn_h     = softmax(scores_h)                          # [1 x (t+1)]
 9.      out_h      = attn_h V_cache[:t+1, h, :]                  # [1 x d_k]
10.  end for
11.  return Concat(out_1, ..., out_H) W^O                          # [1 x d_model]

Walking through line by line:

• Line 2 projects the input into $H$ query vectors. Cheap: $d_{model} \cdot d_k$ FLOPs per head.
• Lines 3–4 compute the new K and V vectors for this token only. They will be appended to the cache.
• Lines 5–6 are writes to the KV cache. In aggregate these grow the cache from $L \cdot H \cdot d_k$ to $(L+1) \cdot H \cdot d_k$ per K, and the same for V.
• Line 7 is the dot product of the current query against all $L+1$ cached keys. This requires reading the entire $K_{cache}$ along the sequence axis for head $h$ — that is, reading $(L+1) \cdot d_k$ values from memory.
• Line 8 is a cheap softmax over a length-$L+1$ vector.
• Line 9 multiplies the attention weights by the cached values, reading $(L+1) \cdot d_k$ values for head $h$.
• Line 11 concatenates and applies $W^O$.

The crucial observation: lines 7 and 9, summed across all $H$ heads, read $2 H (L+1) d_k$ FP16 values from HBM. This is the KV cache traffic per decode step per layer.

3. The KV Cache: Where Inference Memory Goes

The KV cache is the persistent state of the autoregressive decoder. Without it, every decode step would have to recompute $K$ and $V$ from scratch for every previous token, which would scale as $O(L^2 d_{model})$ FLOPs and quickly become intractable. With it, the per-step compute is $O(L d_{model})$ FLOPs but the per-step memory traffic is $O(H L d_k) = O(L d_{model})$ bytes.

Let me put concrete numbers on this. Consider a model with the following dimensions, which roughly corresponds to a 7B LLaMA-style architecture:

• Layers: $N_L = 32$
• Heads: $H = 32$
• Head dim: $d_k = 128$
• Sequence length: $L = 2048$
• Precision: FP16 (2 bytes per value)
• Batch: $B = 1$

The total KV cache size is:

$$S_{KV} = 2 \cdot N_L \cdot H \cdot d_k \cdot L \cdot B \cdot \text{bytes\_per\_value}$$ $$= 2 \cdot 32 \cdot 32 \cdot 128 \cdot 2048 \cdot 1 \cdot 2 = 1{,}073{,}741{,}824 \text{ bytes} = 1 \text{ GiB}$$

For a 70B-scale model ($N_L = 80, H = 64, d_k = 128$) at $L = 4096$: $S_{KV} = 2 \cdot 80 \cdot 64 \cdot 128 \cdot 4096 \cdot 2 = 10.7$ GiB. Just the KV cache for a single request. Now multiply by 32 concurrent users on a serving node and the KV cache becomes the dominant occupant of HBM, even larger than the weights.

The MHA KV cache scales linearly with all of $N_L$, $H$, $d_k$, $L$, and $B$. Two of these — $H$ and $L$ — are exactly what makes long-context, high-throughput serving hard.

flowchart TB
    subgraph HBM["GPU HBM (80 GB on A100)"]
        W[Weights ~14 GB for 7B fp16]
        KV[KV cache for batch=32, L=2048<br/>= 32 x 1 GiB = 32 GiB]
        ACT[Activations small]
    end
    subgraph SM["Streaming Multiprocessors compute"]
        ATTN[Attention head h]
    end
    HBM -- 2 TB/s --> SM
    SM -- write back --> HBM

The HBM-to-SM bandwidth on an A100 is about 2.0 TB/s; on H100 about 3.35 TB/s. The compute throughput is much higher: A100 FP16 tensor cores deliver 312 TFLOP/s. The crossover (where memory bandwidth and compute saturate together) for FP16 matmuls is at arithmetic intensity around 156 FLOP/byte. Decode attention sits far below this crossover.

4. The Memory Bandwidth Bottleneck at Decode Time

Let’s count both FLOPs and bytes per decode step at one MHA layer. With the 7B-style numbers above ($H = 32, d_k = 128, L = 2048, d_{model} = 4096$):

FLOPs per layer per decode step:

• $Q, K, V$ projection: $3 \cdot d_{model} \cdot H d_k = 3 \cdot 4096 \cdot 4096 \approx 50.3$ MFLOPs
• Attention scores: $H \cdot L \cdot d_k = 32 \cdot 2048 \cdot 128 \approx 8.4$ MFLOPs
• Attention values: same, $\approx 8.4$ MFLOPs
• $W^O$: $d_{model}^2 = 16.8$ MFLOPs
• MLP (typical 4x): $\approx 134.2$ MFLOPs

Sum: roughly $\approx 218$ MFLOPs $\approx 2.2 \times 10^8$ FLOPs.

Bytes read per layer per decode step:

• Weights ($Q, K, V, O$ projections plus MLP): $\approx 12 d_{model}^2 \cdot 2 = 12 \cdot 16.8 \text{ M} \cdot 2 \approx 400$ MB just for this layer’s weights. Wait — this is the per-layer weight footprint, around 400 MB at fp16 for a 7B-style block — but it has to be loaded once per decode step.
• KV cache traffic: $2 \cdot H \cdot L \cdot d_k \cdot 2 = 2 \cdot 32 \cdot 2048 \cdot 128 \cdot 2 = 32$ MB
• Activations: small, $\approx 16$ KB.

Total bytes per layer per step: $\approx 432$ MB. Across 32 layers, $\approx 13.8$ GB of memory traffic per token.

Time to read 13.8 GB on A100 (2 TB/s): $13.8 / 2000 \approx 6.9$ ms. Time to do the 32-layer compute ($\sim 7 \times 10^9$ FLOPs) at 312 TFLOP/s: $\sim 0.022$ ms — over 300× faster than the memory traffic.

The implication is stark: at batch 1, the GPU is idle 99.7% of the time waiting for memory. The only fix is to reduce bytes per step. Reducing weights is hard — that’s the model size. Reducing the KV cache is comparatively much easier, because the KV cache has a structural redundancy that we can exploit: many of the $H$ KV head representations are correlated.

flowchart LR
    subgraph Cycle["One decode step"]
        L1[Load weights ~400 MB] --> C1[Compute Q,K,V]
        C1 --> L2[Load KV cache ~32 MB]
        L2 --> C2[Attention]
        C2 --> L3[Load MLP weights ~250 MB]
        L3 --> C3[Compute MLP]
        C3 --> L4[Write KV ~16 KB]
    end
    style L1 fill:#fbb,stroke:#900
    style L2 fill:#fbb,stroke:#900
    style L3 fill:#fbb,stroke:#900
    style C1 fill:#bfb,stroke:#090
    style C2 fill:#bfb,stroke:#090
    style C3 fill:#bfb,stroke:#090

The red boxes are bandwidth-limited; the green boxes are compute. At batch 1, decoding is a sequence of bandwidth boxes with compute boxes hidden behind them. Increasing batch helps amortize the weight reads across requests (weights are shared), but the KV cache is per-request, so it scales linearly with batch. That makes the KV cache the only term that can be attacked structurally.

5. Multi-Query Attention (MQA) — The Aggressive Shortcut

Noam Shazeer in 2019 proposed Multi-Query Attention: a Transformer where each layer has $H$ query heads, but only one key head and only one value head:

$$\mathrm{head}_h^{MQA} = \mathrm{softmax}\!\left(\frac{Q_h K^\top}{\sqrt{d_k}}\right) V$$

where $K = X W^K$ and $V = X W^V$ are shared across all $h$. The KV cache shrinks from $2 H L d_k$ to $2 L d_k$ — an $H$× reduction. On models with $H = 32$ to $64$, this is a 1–2 order of magnitude reduction in the bandwidth-bound term.

The empirical effect is dramatic. Shazeer’s original numbers (and corroborated by this paper at T5 XXL): MQA achieves something like a 6× decode speedup with single-digit-percent quality drop on machine translation. The downside is quality: MQA is “aggressively compressed” — collapsing 32 KV heads into 1 throws away whatever per-head distinctiveness the K/V projections carried, and on benchmarks like summarization the quality loss is measurable.

GQA’s central thesis is that we don’t have to choose between $H$ KV heads (MHA, maximum quality, maximum bandwidth) and $1$ KV head (MQA, minimum quality loss, maximum bandwidth saving). We can pick any integer between, and we should pick the one on the Pareto frontier — empirically, that turns out to be around $G = 8$.

What This Paper Does

The paper has two technical pieces:

Architecture. Introduce Grouped-Query Attention. Partition the $H$ query heads into $G$ groups (typically $H$ divisible by $G$ with each group of size $H/G$). Each group $g$ has one key head $K_g$ and one value head $V_g$, shared across all $H/G$ queries in that group. KV cache is $2 G L d_k$ per layer instead of $2 H L d_k$. Reduction factor: $H/G$. When $G = H$ we get MHA back; when $G = 1$ we get MQA back.

Uptraining recipe. Rather than training a GQA model from scratch (expensive — at the time, T5 XXL pretraining cost ~10000 TPU-days), convert an existing MHA checkpoint to GQA by:

1. For each group $g$, define its KV projections as the mean of the projections of the $H/G$ MHA heads in that group: $W_g^K = \frac{1}{H/G} \sum_{h \in g} W_h^K$, $W_g^V = \frac{1}{H/G} \sum_{h \in g} W_h^V$.
2. Continue training the modified checkpoint for an additional $\alpha = 5\%$ of the original pretraining step count on the original data distribution. Everything else (query and output projections, MLP, embeddings) is initialized from the MHA checkpoint and continues to be trained normally.

The paper shows that this 5% uptraining is enough to recover essentially all the quality of MHA, even though we’ve collapsed 64 KV heads into 8.

Both pieces matter, but in different ways. The architecture is what gets used; the uptraining recipe is what made early adoption cheap. Today’s frontier models (LLaMA 3, Mistral, etc.) train GQA from scratch — they don’t need uptraining. But the uptraining argument is what convinced the field that GQA is “the right structural prior” — if you can mean-pool an MHA model and recover its quality with 5% of the training, it suggests that the extra KV heads of MHA carried mostly redundant information all along.

flowchart LR
    subgraph MHA["MHA: H=8 KV heads"]
        Q1[Q1]-->K1[K1]
        Q2[Q2]-->K2[K2]
        Q3[Q3]-->K3[K3]
        Q4[Q4]-->K4[K4]
        Q5[Q5]-->K5[K5]
        Q6[Q6]-->K6[K6]
        Q7[Q7]-->K7[K7]
        Q8[Q8]-->K8[K8]
    end
    subgraph GQA["GQA-2: G=2 KV heads"]
        GQ1[Q1]-->GK1[K group 1]
        GQ2[Q2]-->GK1
        GQ3[Q3]-->GK1
        GQ4[Q4]-->GK1
        GQ5[Q5]-->GK2[K group 2]
        GQ6[Q6]-->GK2
        GQ7[Q7]-->GK2
        GQ8[Q8]-->GK2
    end
    subgraph MQA["MQA: 1 KV head"]
        MQ1[Q1]-->MK1[shared K]
        MQ2[Q2]-->MK1
        MQ3[Q3]-->MK1
        MQ4[Q4]-->MK1
        MQ5[Q5]-->MK1
        MQ6[Q6]-->MK1
        MQ7[Q7]-->MK1
        MQ8[Q8]-->MK1
    end

Method: Grouped-Query Attention in Detail

The GQA Architecture

Let $H$ be the number of query heads and $G$ a divisor of $H$. Each query head $h \in \{1, \ldots, H\}$ is assigned to group $g(h) = \lceil h / (H/G) \rceil$. The forward pass for layer $\ell$ becomes:

$$Q_h = X W_h^Q \quad \text{for } h = 1, \ldots, H$$ $$K_g = X W_g^K, \quad V_g = X W_g^V \quad \text{for } g = 1, \ldots, G$$ $$\mathrm{head}_h = \mathrm{softmax}\!\left(\frac{Q_h K_{g(h)}^\top}{\sqrt{d_k}}\right) V_{g(h)}$$ $$\mathrm{GQA}(X) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_H) W^O$$

The only change vs. MHA is the number of K and V projections: $G$ instead of $H$. The query projections are unchanged. The output projection $W^O$ is unchanged. The MLP is unchanged. This is one of the architecture’s virtues: it is a minimal surgical edit.

Now the per-decode-step algorithm becomes:

Algorithm 2: Grouped-Query Attention (decode step t, one layer)
Input:
  x_t in R^{1 x d_model}
  W_h^Q for h = 1..H
  W_g^K, W_g^V for g = 1..G
  K_cache in R^{L x G x d_k}
  V_cache in R^{L x G x d_k}
  W^O in R^{(H d_k) x d_model}

 1.  for g = 1..G do
 2.      k_g_new = x_t W_g^K                            # [1 x d_k]
 3.      v_g_new = x_t W_g^V                            # [1 x d_k]
 4.      K_cache[t, g, :] = k_g_new
 5.      V_cache[t, g, :] = v_g_new
 6.  end for
 7.  for h = 1..H do
 8.      q_h = x_t W_h^Q                                 # [1 x d_k]
 9.      g  = ceil(h / (H/G))                            # group lookup
10.      scores_h = q_h (K_cache[:t+1, g, :])^T / sqrt(d_k)
11.      attn_h   = softmax(scores_h)
12.      out_h    = attn_h V_cache[:t+1, g, :]
13.  end for
14.  return Concat(out_1, ..., out_H) W^O

Line by line:

• Lines 1–6 compute and append the $G$ new K and V vectors. This is the only place where the work differs in cardinality from MHA (which had $H$ here, lines 3–6 of Algorithm 1).
• Lines 7–13 compute $H$ attention outputs as before, but now reading from a shared KV slice per group. Multiple query heads in the same group $g$ all index into K_cache[:, g, :]. On the GPU this matters: a well-implemented kernel can fuse the $H/G$ queries in a group and load each cached K vector once to serve all queries that need it, instead of once per query head.

The total KV cache traffic per decode step per layer is now $2 G L d_k$ bytes (FP16: multiply by 2), down from $2 H L d_k$. With $H = 64, G = 8$, this is an 8× reduction.

Formal Definition and KV Cache Savings

Let $T_{KV}^{MHA} = 2 H L d_k \cdot B \cdot \text{bytes}$ be the per-layer per-step KV cache traffic for MHA at batch $B$. For GQA, $T_{KV}^{GQA} = 2 G L d_k \cdot B \cdot \text{bytes}$. The ratio is $H/G$.

A 70B-scale model has $H = 64$. Choosing $G = 8$ gives an 8× reduction in KV cache size and bandwidth:

Model	$H$	$G$	KV cache for $L=4096$, $B=1$	Per-step KV traffic
7B MHA	32	32	1 GiB	32 MB/layer
7B GQA-8	32	8	0.25 GiB	8 MB/layer
70B MHA	64	64	10.7 GiB	64 MB/layer
70B GQA-8	64	8	1.3 GiB	8 MB/layer
70B MQA	64	1	0.17 GiB	1 MB/layer

For long contexts the absolute numbers explode: at $L = 32768$, 70B MHA is $\sim 85$ GiB of KV cache — larger than a single GPU’s HBM. GQA-8 brings this to $\sim 10.7$ GiB, fitting comfortably. The bandwidth savings translate almost linearly into decode wall-clock speedup because the workload is bandwidth-bound.

xychart-beta
    title "KV cache size vs sequence length (70B-scale, fp16)"
    x-axis [1024, 2048, 4096, 8192, 16384, 32768]
    y-axis "KV cache (GiB)" 0 --> 90
    line [2.7, 5.4, 10.7, 21.5, 42.9, 85.9]
    line [0.34, 0.67, 1.34, 2.68, 5.36, 10.7]
    line [0.04, 0.08, 0.17, 0.34, 0.67, 1.34]

The three lines are MHA (top, blue in canonical Mermaid coloring), GQA-8 (middle), MQA (bottom). The gap between MHA and GQA-8 is exactly 8×; between GQA-8 and MQA, exactly 8× again.

The Uptraining Recipe

This is the second contribution and the engineering reason early adopters could afford to switch. Suppose you already trained an MHA model — you spent $C$ TPU-days getting there. You want to ship a faster variant for serving without throwing away $C$ TPU-days of progress.

The uptraining recipe:

Algorithm 3: MHA -> GQA Checkpoint Conversion + Uptraining
Input:
  MHA checkpoint with W_h^K, W_h^V for h = 1..H, plus all other params
  Target group count G (divisor of H)
  Uptraining budget alpha (e.g., 0.05)
  Original pretraining step count T_orig

 1.  group_size = H / G
 2.  for g = 1..G do
 3.      // Mean-pool MHA KV projections within group g
 4.      W_g^K  = (1/group_size) * sum_{h in group(g)} W_h^K
 5.      W_g^V  = (1/group_size) * sum_{h in group(g)} W_h^V
 6.  end for
 7.  Replace W_h^K, W_h^V in checkpoint with the G mean-pooled W_g^K, W_g^V
 8.  Keep all W_h^Q, W^O, MLP, embedding params unchanged
 9.  Continue pretraining for T_uptrain = alpha * T_orig steps
10.  Use the same optimizer (Adafactor), schedule, and corpus as the original run

Line-by-line rationale:

• Lines 1–6. The mean pool is the bit that matters. It is the maximum-likelihood projection of $H/G$ Gaussian-like representations onto a single shared one, assuming equal-variance noise and zero prior preference for any head. Equivalently, in the linear case, the operation $W_g^K = \mathrm{mean}_{h \in g} W_h^K$ is the projection of $H/G$ vectors onto their centroid in parameter space.
• Line 7. This rewrite changes parameter shapes from $(H, d_{model}, d_k)$ to $(G, d_{model}, d_k)$. Optimizer state for those parameters has to be reinitialized for the new shape; the paper notes this is fine because uptraining is short.
• Line 8. Critically, we do not modify the query projections. Each query head keeps its own $W_h^Q$ — that diversity is what we want to preserve. We only collapse on the K and V sides.
• Line 9. $T_{uptrain} = 0.05 \cdot T_{orig}$. For T5 XXL with original $T_{orig} \approx 1$M steps, $T_{uptrain} \approx 50$k steps. The paper reports 600 TPUv3 chip-days for the XXL conversion.
• Line 10. Same optimizer, same LR schedule (which means we restart from the end of the original schedule, i.e., the small final LR), same data. This is critical for not destabilizing the existing weights.

flowchart TB
    MHACkpt[MHA checkpoint] --> Group[Group H heads into G groups]
    Group --> Mean[Mean-pool W_K, W_V within each group]
    Mean --> NewCkpt[New checkpoint with G KV heads]
    NewCkpt --> Uptrain[Uptrain for 5% of original steps]
    Uptrain --> Final[GQA model ready to serve]
    style Mean fill:#bdf,stroke:#039
    style Uptrain fill:#fdb,stroke:#930

Why Mean Pooling? A Principled Analysis

Why pick the mean of the KV projections rather than alternatives? The paper ablates three options:

1. Mean pool (chosen): $W_g^K = \frac{1}{H/G} \sum_{h \in g} W_h^K$
2. First-head select: $W_g^K = W_{h_0}^K$ where $h_0$ is the first head in group $g$
3. Random init: $W_g^K$ sampled from the original initialization distribution

The empirical ranking (paper Table 4 / Figure 4 in the appendix; my paraphrase): mean pool ≫ first-head ≫ random. Intuitively:

Mean pool preserves the maximum amount of pretraining information. If the $H/G$ heads in a group all encode somewhat similar features (which they usually do for nearby heads, especially after pretraining), then their mean is a close approximation to each of them. The model’s downstream layers, which were trained to consume the output of the concatenated heads, will encounter a perturbation that is small for each head and approximately zero on average.

First-head select discards $H/G - 1$ heads worth of learned features. If we pick head 1 to represent the whole group $\{1, 2, 3, 4\}$, we’ve thrown away whatever heads 2, 3, 4 learned. Even if those heads were highly redundant with head 1, the residual differences carried useful information that we are now obligated to relearn.

Random init throws everything away. This is essentially “train GQA from scratch starting from a partial checkpoint with broken K/V layers.” It can recover with enough uptraining, but the budget needed scales much worse with $G$.

Mathematically, we can be more precise. Suppose during MHA training each head’s K projection was $W_h^K = W^*_g + \varepsilon_h$ where $W^*_g$ is the “true” group-level projection and $\varepsilon_h$ is per-head noise with $\mathbb{E}[\varepsilon_h] = 0, \mathrm{Var}[\varepsilon_h] = \sigma^2 I$. Then the mean estimator $\widehat{W}_g^K = \frac{1}{H/G} \sum_h W_h^K$ has variance $\sigma^2 / (H/G) \cdot I$ — it is the minimum-variance unbiased estimator of $W^*_g$ under this noise model. Picking a single head gives variance $\sigma^2 I$, which is $H/G$ times worse. Random init has no relationship to $W^*_g$ at all.

This argument is informal — MHA heads are not literally noisy copies of a shared mean — but the paper’s results match the prediction: mean-pool wins, first-head is second, random is worst, and the gap widens as $H/G$ grows.

flowchart LR
    subgraph Group1["MHA group of 4 heads"]
        H1[W1_K]
        H2[W2_K]
        H3[W3_K]
        H4[W4_K]
    end
    H1 -->|sum| Plus((+))
    H2 -->|sum| Plus
    H3 -->|sum| Plus
    H4 -->|sum| Plus
    Plus -->|divide by 4| Mean[W_g_K mean-pooled]
    style Mean fill:#bdf,stroke:#039

The Group Count G: How to Choose?

GQA gives a knob: $G$ can be anything from $1$ to $H$ (divisors of $H$ are most natural; you can also handle non-divisors with padding). What controls the choice?

KV cache and decode-time bandwidth scale as $G$. Halving $G$ halves the KV cache.

Quality improves monotonically (empirically) with $G$, but with strong diminishing returns. The paper’s data shows the quality vs $G$ curve flattens by $G = 8$ for T5 XXL with $H = 64$.

Decode speed is a function of how bandwidth-bound the system is. At small batch and long context, decoding is heavily bandwidth-bound, so KV reduction translates almost 1:1 to speedup. At large batch, weight reads dominate KV reads, and KV reduction matters less for throughput (but still matters for memory occupancy, which limits how large the batch can grow).

Tensor parallel sharding introduces a sharp constraint: with TP degree $K$, you want $G$ to be divisible by $K$ so that each TP rank can own an integer number of KV heads. MQA ($G = 1$) cannot be sharded — you have to replicate the single KV head across all $K$ ranks, which defeats some of the bandwidth saving. GQA with $G = 8$ on TP=8 lets each rank own exactly one KV head — clean.

The paper concludes that $G = 8$ is a “favorable middle ground” for T5 XXL ($H = 64$). For modern LLMs the same recipe seems to hold: LLaMA 2 70B uses $G = 8$, LLaMA 3 70B uses $G = 8$ (with $H = 64$), Mistral 7B uses $G = 8$ (with $H = 32$).

xychart-beta
    title "Quality vs Speed Pareto (T5 XXL, schematic)"
    x-axis [0.24, 0.28, 0.40, 0.60, 1.00, 1.51]
    y-axis "Avg score" 46 --> 48
    line [46.6, 47.1, 47.2, 47.2, 47.2, 47.2]

The schematic shows MQA at the leftmost bottom point, GQA-8 just to its right and almost at MHA’s quality, and MHA-XXL at the rightmost top. The Pareto-efficient region is the elbow near GQA-8.

GQA in Tensor Parallel Inference

When a 70B model is sharded across, say, 8 GPUs with tensor parallelism, each layer’s matrix multiplications are split across devices. With MHA, the $H$ heads naturally shard along the head dimension: 64 heads / 8 devices = 8 heads per device, each device computing $Q_h, K_h, V_h$ for its 8 heads.

With MQA, there is only $1$ K head and $1$ V head, but $H = 64$ Q heads. The Q heads still shard 8 per device, but K and V must be replicated across all 8 devices. Each device performs the K and V projection from scratch (small work) and stores its own copy of the cache (1× of the single-head cache per device). The KV cache is thus 8× larger in aggregate than the per-device single-KV-head cost would suggest — though still 8× smaller than the MHA aggregate.

With GQA-8 on TP=8, each device owns exactly one of the 8 KV heads plus the 8 corresponding Q heads. KV is fully partitioned with no replication: each device’s KV cache is $1/8$ of the total, and there is no duplication. This is structurally clean and one of the reasons GQA-8 is so popular at production scale.

flowchart TB
    subgraph TPMHA["MHA, TP=8 (64 query and 64 KV heads)"]
        D0M[Device 0: 8Q+8KV]
        D1M[Device 1: 8Q+8KV]
        D7M[... 8 devices ...]
    end
    subgraph TPMQA["MQA, TP=8 (64 query, 1 KV head replicated)"]
        D0Q[Device 0: 8Q + KV copy 1]
        D1Q[Device 1: 8Q + KV copy 2]
        D7Q[... KV replicated 8x ...]
    end
    subgraph TPGQA["GQA-8, TP=8 (64 query, 8 KV heads sharded)"]
        D0G[Device 0: 8Q + 1KV]
        D1G[Device 1: 8Q + 1KV]
        D7G[... 8 devices ...]
    end
    style TPGQA fill:#dfd,stroke:#090

The takeaway: $G$ should typically be a multiple of (or equal to) the tensor parallel degree $K$. $G = K$ is the cleanest case. $G = 8$ with TP $\in \{1, 2, 4, 8\}$ all work nicely.

Experiments

Models and Datasets

The paper evaluates on T5 (encoder-decoder), specifically T5 Large and T5 XXL with the v1.1 architecture. The base configuration is:

• T5 Large: $\sim 770$M parameters, $H = 16$, $d_k = 64$, $N_L = 24$ (encoder + decoder)
• T5 XXL: $\sim 11$B parameters, $H = 64$, $d_k = 64$ (note: T5 actually uses $d_k = 64$ even at XXL; this is a deliberate T5 choice), $N_L = 24$ (encoder + decoder)

The conversion is applied only to decoder self-attention and encoder-decoder cross-attention. Encoder self-attention is left as MHA — the encoder is not bandwidth-bound (it processes the input in parallel and is compute-bound), so changing it would only slow things down without saving anything.

Datasets:

• CNN/DailyMail (summarization, news)
• arXiv (summarization, long scientific papers)
• PubMed (summarization, biomedical abstracts)
• MediaSum (summarization, dialogue)
• MultiNews (summarization, multi-document)
• WMT EN-DE (machine translation)
• TriviaQA (open-domain question answering)

This dataset selection is deliberate: long-input summarization tasks stress the KV cache the hardest, because the decoder must attend to many encoder tokens (cross-attention K and V are long), and arXiv/PubMed inputs are routinely thousands of tokens. If GQA were going to fail anywhere, it would be on these long-context summarization tasks.

Inference Timing Methodology

The paper measures inference time per sample on a single TPU. The measurements are decode-only — i.e., the wall-clock time for the autoregressive part of generation, not counting prefill. This is the right thing to measure because, as established above, decode is the bandwidth-bound phase.

Importantly: timing is done with greedy decoding (no sampling overhead) and a batch size of 1 (the worst case for bandwidth bottleneck, and also the most representative for low-latency single-user serving).

Main Results

Here is the canonical table the paper produces for T5 XXL with 5% uptraining:

Model	$T_{infer}$ (s/sample)	Avg	CNN/DM	arXiv	PubMed	MediaSum	MultiNews	WMT EN-DE	TriviaQA
MHA-Large	0.37	46.0	42.9	44.6	46.2	35.5	46.6	27.7	78.2
MHA-XXL	1.51	47.2	43.8	45.6	47.5	36.4	46.9	28.4	81.9
MQA-XXL	0.24	46.6	43.0	45.0	46.9	36.1	46.5	28.5	81.3
GQA-8-XXL	0.28	47.1	43.5	45.4	47.7	36.3	47.2	28.4	81.6

Several quantitative observations:

• MHA-XXL → MQA-XXL: 6.3× speedup (1.51 → 0.24 s), 0.6 average points of quality lost.
• MHA-XXL → GQA-8-XXL: 5.4× speedup (1.51 → 0.28 s), only 0.1 average points lost.
• GQA-8-XXL vs MHA-Large: GQA-8 is XXL-scale (with XXL-scale quality) but runs at faster than Large-scale speed (0.28 vs 0.37 s). This is the headline result: GQA-8 gives you XXL-class quality at sub-Large-class latency.
• PubMed and MultiNews even exceed MHA-XXL with GQA-8 (47.7 vs 47.5; 47.2 vs 46.9). This is mild and likely within noise, but it indicates GQA is not in any meaningful sense a “weaker” architecture — the implicit regularization of fewer KV heads may even help slightly on some tasks.

xychart-beta
    title "T5 XXL: quality vs decode time"
    x-axis [0.20, 0.30, 0.40, 0.60, 1.00, 1.51]
    y-axis "Average score" 45.5 --> 47.5
    line [46.6, 47.1, 47.2, 47.2, 47.2, 47.2]

Ablation 1: Checkpoint Conversion Strategy

The paper ablates the initialization strategy for the KV projections during conversion. Three methods compared:

1. Mean pool: $W_g^K = \frac{1}{H/G} \sum_h W_h^K$
2. First head: $W_g^K = W_1^K$ (just pick the first head from each group)
3. Random: re-initialize $W_g^K$ from scratch

Result: mean pool > first head > random by a comfortable margin. At very small uptraining budgets (1% of original), the gap is large; at large budgets (5%+), random eventually catches up but never quite reaches mean-pool quality.

This is the paper’s strongest evidence that pretraining has already learned approximately-shared structure within each group, and mean-pooling is the right way to extract that shared structure for GQA.

Ablation 2: Uptraining Budget

The paper sweeps the uptraining budget $\alpha \in \{0, 1\%, 5\%, 10\%\}$. The key qualitative findings:

• $\alpha = 0$: no uptraining at all. GQA is surprisingly close to MHA already — the mean-pooled checkpoint is a decent initialization just by itself. MQA, however, is much worse without uptraining.
• $\alpha = 1\%$: most of the gap is closed.
• $\alpha = 5\%$: essentially saturated for GQA. MQA is also close but slightly worse.
• $\alpha = 10\%$: marginal further gains.

The paper uses $\alpha = 5\%$ as the default, which is small enough to be affordable (you spent 100 days pretraining; 5 more days to convert) and large enough to capture nearly all the recoverable quality.

xychart-beta
    title "Schematic: quality vs uptraining alpha"
    x-axis [0, 1, 5, 10]
    y-axis "Avg score" 44 --> 48
    line [46.2, 46.8, 47.1, 47.15]
    line [44.5, 45.8, 46.6, 46.7]

The upper curve is GQA-8 (rapid recovery, plateaus at MHA quality). The lower curve is MQA (slower recovery, plateaus below MHA quality).

Ablation 3: Number of Groups

Sweeping $G$ from 1 (MQA) to 64 (MHA) at T5 XXL:

$G$	Speed	Quality
1 (MQA)	very fast	-0.6 vs MHA
2	fast	-0.3
4	fast	-0.15
8 (GQA-8)	fast	-0.1
16	medium	-0.05
32	slow	-0.02
64 (MHA)	slowest	0 (reference)

The quality curve is concave — most of the gap closes by $G = 8$, then long flat tail. The speed curve is roughly linear in $G$ in the bandwidth-bound regime: each doubling of $G$ doubles the KV traffic.

The intersection point — where additional $G$ buys progressively less quality but progressively more cost — is $G = 8$ for T5 XXL with $H = 64$.

A Closer Look at the Cost Accounting

Before moving on to interpretation, let me make the cost accounting fully explicit. This is the kind of arithmetic that I find clarifying and that the paper itself only sketches.

Consider one decoder layer of a 70B-scale model at decode time, with $L = 8192$ context and batch $B = 1$.

MHA case ($H = 64$, $d_k = 128$, $d_{model} = 8192$):

• Q projection: $1 \cdot d_{model} \cdot H \cdot d_k = 1 \cdot 8192 \cdot 64 \cdot 128 \approx 67$ MFLOPs.
• K projection: same, $\approx 67$ MFLOPs.
• V projection: same, $\approx 67$ MFLOPs.
• Attention scores ($Q K^\top$): $H \cdot L \cdot d_k = 64 \cdot 8192 \cdot 128 \approx 67$ MFLOPs.
• Attention output ($\mathrm{softmax} \cdot V$): same, $\approx 67$ MFLOPs.
• $W^O$: $1 \cdot d_{model} \cdot d_{model} \approx 67$ MFLOPs.
• MLP (typical 4× expansion): $\approx 540$ MFLOPs.

Total compute per layer per step: $\approx 942$ MFLOPs $\approx 10^9$ FLOPs.

Now bytes:

• Weights for $Q, K, V$: $3 \cdot d_{model} \cdot H \cdot d_k \cdot 2 = 3 \cdot 8192 \cdot 8192 \cdot 2 = 400$ MB.
• Weights for $W^O$: $d_{model}^2 \cdot 2 = 134$ MB.
• Weights for MLP: $\approx 1.07$ GB.
• KV cache traffic: $2 \cdot H \cdot L \cdot d_k \cdot 2 = 2 \cdot 64 \cdot 8192 \cdot 128 \cdot 2 = 268$ MB.
• Activations: negligible.

Total bytes: $\approx 1.87$ GB per layer per decode step. At 80 layers, that’s roughly 150 GB of memory traffic per decode token, which on an H100 (3.35 TB/s) takes $\sim 45$ ms.

GQA-8 case (same model, $G = 8$ instead of $H = 64$ KV heads):

• Q projection unchanged: $\approx 67$ MFLOPs, 134 MB of weight reads (since Q weight matrix is $d_{model} \cdot H \cdot d_k = d_{model}^2$).
• K projection: $1 \cdot d_{model} \cdot G \cdot d_k = 1 \cdot 8192 \cdot 8 \cdot 128 \approx 8.4$ MFLOPs. Weight: $d_{model} \cdot G \cdot d_k \cdot 2 = 16.8$ MB.
• V projection: same as K, $\approx 8.4$ MFLOPs, 16.8 MB.
• Attention scores: $H \cdot L \cdot d_k = 67$ MFLOPs (queries still split across $H = 64$).
• Attention output: same.
• $W^O$: $134$ MB weight (unchanged).
• MLP: 1.07 GB weight (unchanged).
• KV cache traffic: $2 \cdot G \cdot L \cdot d_k \cdot 2 = 2 \cdot 8 \cdot 8192 \cdot 128 \cdot 2 = 33.5$ MB.

Total bytes per layer per step: $134 + 16.8 + 16.8 + 134 + 1070 + 33.5 \approx 1.40$ GB. Across 80 layers, $\approx 112$ GB per token. At H100 bandwidth, $\sim 33$ ms.

Speedup: $45 / 33 \approx 1.36$×. That’s surprisingly modest given KV traffic dropped 8×. The reason: at $L = 8192$ on this 70B model, KV is only $\sim 14\%$ of bytes; reducing it to $\sim 2\%$ saves at most that 14% directly, plus a sliver from smaller K/V projection weight reads.

Now repeat at $L = 32768$:

• MHA KV traffic: $4 \times$ the $L = 8192$ figure = $\approx 1.07$ GB per layer.
• Total MHA bytes per layer: $\approx 2.67$ GB. Across 80 layers: 214 GB. At H100 BW: 64 ms.
• GQA-8 KV traffic: $\approx 134$ MB per layer.
• Total GQA-8 bytes per layer: $\approx 1.50$ GB. Across 80 layers: 120 GB. At H100 BW: 36 ms.

Speedup: $64 / 36 \approx 1.78$×. The advantage grows with context length, exactly as the linear scaling predicts.

Even longer: $L = 131072$:

• MHA KV traffic: $16 \times$ the $L = 8192$ figure $\approx 4.29$ GB per layer.
• Total MHA bytes per layer: $\approx 5.59$ GB. Across 80 layers: 447 GB. At H100 BW: 133 ms per token.
• GQA-8 KV traffic: $\approx 537$ MB per layer.
• Total GQA-8 bytes per layer: $\approx 1.90$ GB. Across 80 layers: 152 GB. At H100 BW: 45 ms per token.

Speedup: $133 / 45 \approx 2.95$×.

At very long contexts, GQA approaches a full $H/G = 8$× speedup because KV completely dominates the byte budget. At short contexts, GQA gives much smaller speedups because KV is already a small fraction of the budget. The benefit of GQA scales with context length — which is why it has become indispensable as production context windows grew from 2k to 128k and beyond.

xychart-beta
    title "Decode time per token vs context (70B, H100)"
    x-axis [2048, 8192, 32768, 131072]
    y-axis "ms per token" 0 --> 140
    line [33, 45, 64, 133]
    line [32, 33, 36, 45]

The top curve is MHA. The bottom curve is GQA-8. The gap widens with context length.

Analysis: Reading the Numbers

A few things deserve emphasis beyond the raw table.

The 6.3× MQA speedup is not 64×. Even though MQA reduces KV cache by 64× (since $H = 64$), the decode wall-clock speedup is only 6.3×. The reason is that the KV cache traffic is one component of the per-step memory budget. Weights are still loaded every step; activations are still written. Reducing KV traffic from say 60% of bytes to 5% of bytes gives roughly a 2.4× speedup ratio at best on memory time; the rest of the 6.3× comes from removing the K and V projection compute (which, while small, dominates after KV traffic shrinks) and from kernel-level effects like reduced register pressure and improved cache locality.

Why is GQA-8 only marginally slower than MQA? MQA is $G = 1$. GQA-8 has 8× more KV traffic than MQA. Yet the wall-clock difference is 0.28 vs 0.24 s — about 17%, not 700%. This is because at $G = 8$ the KV traffic has already been driven well below the weight traffic — additional reductions only chip away at a small remaining piece of the per-step memory cost. We are now solidly in the regime where weights, not KV, are the bottleneck. This is a key piece of intuition: GQA’s job is to push the KV term below the weight term; once there, additional reduction has rapidly diminishing returns.

Why is uptraining only 5% sufficient? Because mean-pooling is a remarkably good initializer. The mean-pooled KV projection, applied to the input $x_t$, produces $K_g x_t = \frac{1}{H/G} \sum_h K_h x_t$ — the average of the $H/G$ heads’ K outputs. The model’s downstream attention layers, which were trained on the concatenated outputs of $H$ heads, now see attention values that are pooled within each group. Mean-pooling preserves the first moment of each group’s K and V outputs exactly. The downstream layers’ learned linear combinations (the $W^O$ output projection) take linear combinations of attention outputs — to first order, those linear combinations are unchanged when we mean-pool. Only the higher-order structure (variance, individual-head idiosyncrasies) is lost, and that’s what the 5% uptraining patches up.

What scales the recipe? For larger models, two competing effects matter. (a) Larger models have more $d_{model}$, so per-step weight reads scale faster than KV reads scale (the latter scale only as $d_k = d_{model}/H$). KV becomes a smaller fraction of total bandwidth as $d$ grows. (b) Larger models tend to use longer sequences, which makes KV larger in absolute terms. The paper argues (and the LLaMA-3 era confirms) that the net effect is roughly scale-invariant: GQA-8 remains a good choice at every scale tested.

Limitations and Boundary Conditions

GQA is not a free lunch. Several limitations apply:

1. Quality gap is small but not zero. The 0.1-point average drop is well within noise for any single dataset, but it is consistent across many datasets and is therefore likely a real, if small, regression. If you are working in a domain where the last 0.1 of quality matters more than 5× decode speed, stay with MHA.

2. The recipe is specific to the regime studied. T5 is encoder-decoder; the paper applies GQA to decoder self-attention and cross-attention only. For decoder-only models (the dominant paradigm now), the recipe transfers directly to decoder self-attention. The cross-attention case may differ in models where the cross-attention K/V are long (very long encoder inputs).

3. Uptraining requires the original training stack. You need the original optimizer state, data loader, and schedule. If you only have weights and no infra, you have a harder time. (Modern frontier model labs work around this by training GQA from scratch.)

4. KV cache is not the only bottleneck. Once GQA reduces KV traffic below weight traffic, further reductions are wasted. This implies GQA’s benefit is largest for small models with long contexts (KV-dominated) and smallest for large models with short contexts (weight-dominated).

5. Group definitions are fixed. The paper uses consecutive grouping (heads 1..H/G in group 1, etc.). It does not explore learned groupings, which could in principle do better. This is a natural research direction that subsequent work has not heavily pursued — empirically, the choice of grouping seems to matter much less than the choice of $G$.

6. GQA assumes K and V should be tied. Each group has one $K_g$ AND one $V_g$. The paper does not ablate having, e.g., $G_K = 8$ but $G_V = 16$. Subsequent work (e.g., DeepSeek’s MLA) has explored asymmetric KV compression.

Critical Assessment: Weaknesses & Improvements

The preceding sections summarize what the paper claims and what it demonstrates well. This section reads the paper against the grain — where the evidence is weaker than the headline numbers suggest, what the authors do not test, and what a stronger version of this work would look like.

Weaknesses and flaws in the evaluation

The entire empirical case rests on one architecture family (T5) and one task family (conditional generation with a long encoder input). Every number in the main results table — the 47.2 / 47.1 / 46.6 average scores, the 1.51 / 0.28 / 0.24 second timings — comes from T5 Large and T5 XXL, an encoder-decoder architecture that was already niche relative to decoder-only GPT-style models by the time the paper was published in 2023. GQA is now used almost exclusively in decoder-only models (LLaMA, Mistral, Qwen), yet the paper contains zero decoder-only experiments. The generalization from encoder-decoder cross-attention (where the K/V sequence is the encoder output, fixed in length and computed once) to decoder-only self-attention (where the K/V sequence grows with every generated token and is recomputed incrementally) is treated as obvious in the paper, but the two settings stress the KV cache very differently — cross-attention K/V never grows during decode, so the argument that “GQA saves bandwidth because the KV cache keeps growing” does not even apply to half of the paper’s own experimental setting (cross-attention). The self-attention numbers that would have mattered most to today’s readers are not reported at all; they were reconstructed years later by the community (LLaMA 2, Mistral) rather than by the authors.

No confidence intervals, variance, or seeds are reported anywhere in the paper. Every number in Table 1 (the 47.2/47.1/46.6 comparison) is a single point estimate from what appears to be a single training/uptraining run per configuration. The paper’s own headline claim — “GQA-8 loses only 0.1 points versus MHA” — is a difference of 0.1 on a metric that averages seven different task scores together; there is no report of run-to-run variance for either MHA or GQA-8 checkpoints, so it is not possible to tell from the paper alone whether 0.1 points is inside or outside the noise band of a single training run. Given that uptraining runs for only 5% of pretraining steps and starts from a perturbed checkpoint, some run-to-run variance is almost certain, and the paper’s central selling point (“nearly free quality”) would be considerably weakened if that variance turned out to be, say, ±0.2 points.

The ablation on group count $G$ (Table/Figure with $G \in \{1,2,4,8,16,32,64\}$) is run only at one model scale (T5 XXL) and only on the averaged score, not per-task. It’s entirely possible that some individual tasks (e.g., TriviaQA, which stresses factual recall more than summarization) have a different, less forgiving $G$-vs-quality curve than the average suggests. Averaging across seven very different tasks (summarization, translation, QA) before reporting the ablation numbers can hide task-specific cliffs — a phenomenon well documented in later compression literature (e.g., SliceGPT, GQA follow-ups) where per-task degradation is far less uniform than an average score implies. Without per-task breakdowns for the $G$-sweep specifically (the paper only gives per-task breakdowns for the headline MHA/MQA/GQA-8 comparison in Table 1, not for the full $G$ sweep), a practitioner cannot know whether their specific downstream task sits near the “cliff” or safely on the flat part of the curve.

The paper never controls for total training compute across the MHA, MQA, and GQA-8 conditions. GQA-8-XXL and MQA-XXL both receive the same uptraining budget ($\alpha = 5\%$) starting from the same MHA-XXL checkpoint, but MHA-XXL itself received 100% of the original pretraining budget with none held back for “free” continued training. This means the comparison in Table 1 is not “three checkpoints trained under equal compute,” it’s “one checkpoint trained for $T$ steps vs. two checkpoints trained for $T + 0.05T$ steps.” The extra 5% of gradient steps is not free — it’s more total training compute than the MHA-XXL baseline received. A cleaner ablation would continue training MHA-XXL for the same additional 5% of steps (with no architecture change) and show that this extra training alone doesn’t already close most of the gap. The paper does not report this control, so a skeptical reader cannot fully separate “GQA architecture is good” from “5% more training helps everything a little.”

Limitations the authors understate or omit

The paper is silent on training-from-scratch GQA vs. uptrained GQA. All of the paper’s GQA results come from uptraining an MHA checkpoint. The paper does not report what happens if you train a GQA model from scratch for the full pretraining budget (not just uptraining from an MHA initialization). This matters because uptraining implicitly benefits from the MHA checkpoint’s already-learned representations; a from-scratch GQA model has to learn the shared group-level KV representations without that head start. The paper’s own framing (“GQA gives nearly all of MHA’s quality”) is therefore only rigorously established for the uptrained regime, yet by the time GQA became the industry default (LLaMA 3, Mistral), essentially every production deployment trains GQA from scratch. The paper’s central empirical claim and the way GQA is actually used in the field are subtly mismatched, and the paper does not flag this gap explicitly.

The tensor-parallelism argument for $G=8$ is asserted, not measured. Section on TP sharding claims that $G$ being a multiple of the TP degree $K$ avoids KV replication overhead, but the paper reports no wall-clock or memory measurements for TP-sharded MQA vs. GQA-8 vs. MHA. The entire TP-efficiency argument in the paper (and repeated in most GQA explainers, including earlier in this review) is a plausibility argument backed by first-principles reasoning, not an experiment. Given how much production motivation for $G=8$ specifically (rather than $G=4$ or $G=16$) rests on this TP-sharding cleanliness, the absence of a direct measurement is a real gap.

The uptraining budget sweep ($\alpha \in \{0,1,5,10\%\}$) stops at 10%, right where the curve is still (barely) rising for MQA. The paper picks $\alpha=5\%$ as “the” default without testing whether $\alpha=15\%$ or $\alpha=20\%$ would let MQA fully close the gap to MHA — which would undercut the paper’s implicit argument that MQA has an irreducible quality ceiling that GQA avoids. It’s possible MQA simply needs more uptraining rather than more KV heads; the paper’s experimental design (stopping the sweep at 10%) cannot distinguish these two hypotheses.

The paper does not discuss what happens under batched, high-throughput serving — the now-dominant real-world deployment regime. All timing numbers are batch-size-1 decode latency on a single TPU. As this review’s own “cost accounting” section shows, KV cache’s share of the total memory budget (and hence GQA’s speedup) depends heavily on batch size, sequence length, and hardware — none of which the paper varies. The batch-1 numbers are the most flattering possible case for GQA (KV is a larger fraction of bytes moved when other costs like weight-reads are amortized less); at high batch sizes where weight reads are amortized across many requests, GQA’s relative benefit could look rather different, and the paper offers no data point there.

Concrete improvement suggestions

1. Report per-run variance. Rerun the T5 XXL MHA/MQA/GQA-8 training and uptraining with at least 3 different random seeds and report mean ± std for the average score. Without this, the paper’s central quantitative claim (“only 0.1 points lost”) cannot be distinguished from noise.
2. Add a compute-matched MHA control. Continue training the MHA-XXL baseline for the same extra 5% of steps with no architecture change, and report whether that alone recovers some of the apparent GQA/MQA quality advantage over pure-MHA-at-$T_{orig}$ steps. This is the single most important missing ablation in the paper.
3. Run at least one decoder-only architecture end to end. Even a modest decoder-only model (GPT-2 scale would have been feasible in 2023) trained and uptrained under the same protocol would directly validate the paper’s implicit claim that the cross-attention findings transfer to self-attention — the setting that ended up mattering for the entire field.
4. Measure TP-sharded throughput directly, rather than arguing from first principles that $G$ divisible by TP degree $K$ avoids replication overhead. A simple experiment sharding MQA, GQA-4, GQA-8, and MHA across TP=2/4/8 and reporting measured throughput and memory would settle the question the paper only asserts.
5. Extend the uptraining-budget sweep past 10% specifically for MQA, to determine whether MQA’s quality ceiling is a genuine architectural limit (fewer KV heads can’t represent enough information no matter how much you train) or merely a slower convergence rate that more uptraining would eventually close. This distinguishes “MQA is architecturally worse” from “MQA just needs a bigger uptraining budget,” which the current experimental design cannot separate.

Reproducibility and Practical Notes

For practitioners wanting to implement GQA today:

• PyTorch: since version 2.0, F.scaled_dot_product_attention accepts queries of shape $[B, H_q, L, d_k]$ and keys/values of shape $[B, H_{kv}, L, d_k]$ with $H_q$ a multiple of $H_{kv}$, and handles broadcasting natively. FlashAttention 2 has native GQA support.
• vLLM, TGI, TensorRT-LLM: all major serving stacks support GQA out of the box. Setting the model config’s num_key_value_heads = G is typically the only knob.
• Conversion script: if you have an MHA checkpoint and want to convert to GQA-$G$, the operation in PyTorch is approximately:

W_K_mha = state_dict["k_proj.weight"]            # [H * d_k, d_model]
W_K_mha = W_K_mha.view(H, d_k, d_model)          # [H, d_k, d_model]
W_K_gqa = W_K_mha.view(G, H // G, d_k, d_model).mean(dim=1)
state_dict["k_proj.weight"] = W_K_gqa.reshape(G * d_k, d_model)

Repeat for $W^V$. Continue training as usual.

• Choosing $G$. Default to $G = 8$ if your $H \geq 16$. If you’re tensor-parallel sharding across $K$ devices, ensure $G$ is a multiple of $K$. Common production choices:
  • LLaMA-2 70B / LLaMA-3 70B: $H = 64, G = 8$
  • Mistral 7B: $H = 32, G = 8$
  • LLaMA-3 8B: $H = 32, G = 8$
  • Qwen2.5 72B: $H = 64, G = 8$
• Long contexts: GQA’s importance grows linearly with context length. At 128k context, the KV cache for an MHA 70B model is impractical; GQA-8 makes it feasible on commodity hardware.

A few additional engineering notes that the paper does not cover explicitly but are important in practice:

KV cache layout. The natural layout for an MHA cache is $[L, H, d_k]$ or $[B, H, L, d_k]$. For GQA, the natural layout is $[L, G, d_k]$ (or with batch). The Q tensor remains $[B, H, L, d_k]$ but at decode time has only $L = 1$. The kernel has to broadcast the $G$ KV heads to the $H$ Q heads, which on modern attention kernels (FlashAttention-2’s GQA support, xformers) is handled by repeating along the head dimension in registers or in the loop tiling, not by materializing a $[B, H, L, d_k]$ KV tensor — that would defeat the bandwidth savings.

Mixed-precision considerations. KV cache is often stored in FP8 or even INT4 in modern serving systems. GQA stacks with this: with $G = 8$ and FP8 KV, the per-step KV bandwidth is reduced by another 2× relative to FP16 GQA-8. The combinations are multiplicative: GQA-8 + FP8 = $8 \cdot 2 = 16$× reduction vs MHA + FP16.

Speculative decoding. When verifying $k$ speculative tokens in a single forward pass, the model loads the KV cache once and attends $k$ new queries against it. The bandwidth cost is dominated by reading the existing $L$ cached entries, and the saving from GQA persists. In fact, GQA increases the speculative-decoding “acceptance rate” headroom because it makes the target model’s verification pass cheaper per token attempted.

Training cost. Training a GQA model from scratch is cheaper than training the equivalent MHA model. Fewer K and V parameters means slightly fewer FLOPs per step (the K and V projections shrink by $H/G$). For a 70B model with $H = 64, G = 8$, the per-step training FLOPs drop by roughly 3% — small but not zero.

Parameter savings. GQA reduces the parameter count slightly. For a layer with $d_{model}$ and $H, G$: MHA K+V parameters are $2 \cdot d_{model} \cdot H \cdot d_k$; GQA K+V parameters are $2 \cdot d_{model} \cdot G \cdot d_k$. For LLaMA-2 70B ($d_{model} = 8192, H = 64, G = 8, d_k = 128$), this saves $2 \cdot 8192 \cdot (64 - 8) \cdot 128 = 117$M parameters per layer, times 80 layers = 9.4B parameters. That is large in absolute terms but small relative to the 70B total — most of the parameter budget is in MLPs, not attention.

Conversion practicalities. When converting an HF Transformers MHA checkpoint to GQA, in addition to mean-pooling $W^K$ and $W^V$, you must also update the model config (num_key_value_heads = G) so that the runtime knows to broadcast appropriately. Many model architectures share a common config field for this (LLaMA, Mistral, Mixtral, Qwen all expose num_key_value_heads). Older codebases sometimes call it n_kv_heads or similar.

Numerical effects. Mean-pooling produces $W_g^K$ vectors whose norm is smaller than the original per-head $W_h^K$ norms by roughly $\sqrt{H/G}$ (under independence assumptions). This shifts the attention scores’ magnitudes downward, which after softmax produces flatter attention distributions. In practice this effect is small and the 5% uptraining absorbs it; but for very aggressive conversions ($G$ very small, no uptraining), the score magnitude shift can affect generation quality in subtle ways (more uniform attention → over-smoothed outputs).

Connections to Subsequent Work

GQA was published in mid-2023, and the field has not stood still. A short tour of what followed clarifies where GQA fits in the larger story.

LLaMA 2 (July 2023). Meta’s open-weight model adopted GQA for the 34B and 70B variants, with $G = 8$. This was the first major production deployment and helped establish the recipe as the default. The 7B and 13B variants used MHA, which in retrospect was a missed opportunity — at low scale, GQA is also a clear win.

Mistral 7B (October 2023). Used GQA with $G = 8$ on $H = 32$ — i.e., a smaller $H/G = 4$ ratio. Mistral combined GQA with sliding window attention (each token attends to the previous $W = 4096$ tokens), which further reduces KV cache size to $O(W)$ rather than $O(L)$. The combination is synergistic: GQA reduces the KV traffic per position; sliding window reduces the number of positions.

LLaMA 3 (April 2024). Adopted GQA across all sizes (8B, 70B, 405B). The 8B uses $H = 32, G = 8$; the 70B uses $H = 64, G = 8$. The 405B continues this pattern. The constant $G = 8$ across scales is the empirical sweet spot the field has settled on.

DeepSeek-V2 / V3 with Multi-head Latent Attention (May 2024 onwards). MLA is the next-generation idea: instead of grouping $H$ KV heads into $G$ KV heads, compress KV into a single low-rank latent $c_t \in \mathbb{R}^{d_c}$ with $d_c \ll H d_k$. At decode time only $c_t$ is cached; the per-head K and V are reconstructed on the fly via small matrix multiplications. This is a more aggressive structural prior than GQA — empirically, DeepSeek-V2 reports MLA outperforming GQA at the same KV-cache budget. The trade is more compute per step but less memory traffic; on memory-bound decode this is favourable.

Speculative decoding. GQA composes cleanly with speculative decoding (where a small draft model proposes tokens that a large target model verifies). The bottleneck for the target model’s verification step is still KV cache bandwidth, and GQA helps the verifier the same way it helps a normal decoder.

Paged attention. vLLM’s PagedAttention manages KV cache in fixed-size blocks for better memory utilization. PagedAttention works at the level of “one block per KV head”; with GQA’s $G$ KV heads, blocks are smaller and packing efficiency improves. Effectively GQA also helps the memory allocator.

The picture is that GQA is a primitive that compounds with every other inference optimization. This is what makes it durable — not that it’s the final word on attention efficiency, but that it sits at the right level of abstraction for everything else to build on.

My Take

GQA is one of the cleanest “small idea, large impact” papers in recent LLM literature. The core insight — that you don’t have to pick between $H$ KV heads and $1$ KV head, you can pick any divisor in between — is in retrospect obvious. The technical depth is in the uptraining recipe, particularly the mean-pooling argument, which gives a principled way to convert existing MHA checkpoints cheaply.

Why it has had outsized impact:

• It’s a minimal architectural change. No new operators, no new activations, no new optimizers. Just change the shape of two weight matrices per layer. This makes it trivial to retrofit into existing codebases.
• It cleanly composes with tensor parallelism. $G$ heads can be sharded across $G$ TP ranks with no replication, no all-reduce overhead, no awkward edge cases.
• It composes with everything else. GQA + Flash Attention: works. GQA + paged KV cache (vLLM): works. GQA + speculative decoding: works. GQA + sliding window attention (Mistral): works. GQA + quantized KV: works.
• It hits a sweet spot in the design space. $G = 8$ at $H = 64$ recovers nearly all the quality of MHA at nearly all the speed of MQA. The fact that “8” is the magic number across many model scales suggests it reflects a structural property of attention rather than a tuned hyperparameter.

What the paper doesn’t fully resolve, and where the field has gone since:

• Asymmetric KV. Why should $K$ and $V$ have the same group count? DeepSeek-V2’s Multi-head Latent Attention (MLA) compresses KV jointly with a low-rank latent, which is a much more aggressive reduction than GQA and arguably the next step on the same trajectory.
• Learned grouping. No prior work has convincingly shown that learned head-to-group assignments beat the trivial consecutive assignment. This is interesting because it suggests heads are largely interchangeable at the group-level granularity.
• Beyond pretraining recovery. The mean-pooling argument is most compelling when the KV heads within a group were already correlated. For models trained from scratch with GQA, this argument doesn’t apply directly. Yet GQA still works. The deeper question — why is the marginal value of KV heads so low — is not fully answered.

The empirical fact that GQA is essentially free at $G = 8$ has been one of the quiet revolutions of LLM serving. It is the reason long-context serving (32k, 128k) is feasible at the price points it is. Combined with FlashAttention, paged KV, and speculative decoding, GQA is one of the four pillars of modern LLM inference. This paper is short, well-targeted, and has aged exceptionally well.

References

• 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. EMNLP 2023. arXiv:2305.13245.
• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150.
• Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS 2017.
• Raffel, C. et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM). SOSP 2023.
• Touvron, H. et al. (2023). LLaMA 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288.
• DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.
• Jiang, A. et al. (2023). Mistral 7B. arXiv:2310.06825.
• Pope, R. et al. (2022). Efficiently Scaling Transformer Inference. MLSys 2023. (The PaLM inference paper — establishes the bandwidth-bound decode framework GQA builds on.)