Core ML

Grouped-Query Attention

Instead of every head keeping its own keys and values or all heads sharing one pair, grouped-query attention splits heads into a handful of groups where each group shares a KV pair — most of the cache saving, almost none of the quality loss.

Grouped-query attention sits between MHA and MQA: four query heads share two KV groups, so each group serves two heads, cutting the KV cache in half while keeping more representational diversity than a single shared pair

Why Does This Exist?

Multi-query attention solved the KV cache problem aggressively: collapse all heads to one shared KV pair and you get an H-fold reduction in cache memory. The problem was quality. Models trained with a single shared KV pair showed measurable degradation on tasks that relied on heads attending from different angles — translation, long-document QA, coding tasks where context structure matters.

Ainslie et al. at Google (2023) asked whether you actually need to go all the way to one group. They ran a systematic ablation: train the same model with G groups where G ranged from 1 (MQA) to H (full MHA), and measure the quality-cache trade-off at each point. The answer was that G=8 or so recovered nearly all of MHA's quality while cutting the KV cache by 4× compared to MHA. G=1 (MQA) saved more cache but paid meaningfully more in quality. The sweet spot existed, and it wasn't at either extreme.

That finding shipped as Llama 2 70B (8 KV groups out of 64 query heads) and Mistral 7B (8 groups out of 32), among others. Grouped-query attention is now the default in virtually every open-weight model that runs in production.

Think of It Like This

Table groups at a conference

Picture 32 researchers at a conference, seated at 8 tables. Every researcher at a table shares the same reference binder — the keys and values their group uses. Each researcher still writes their own question (query), and each table's binder is slightly different from the others. The researchers at table 1 have one binder; table 2 has its own; but within a table, the binder is shared.

This is grouped-query attention. Eight different binders instead of 32 separate ones (MHA) or one shared binder for the whole room (MQA). The reduction in binders translates directly to a reduction in the KV cache.

How It Actually Works

Projections and grouping

In GQA, query heads are divided into G groups. Each group gets one K projection and one V projection, each of shape $[d_{model}, d_{head}]$ . The query projection stays full: H separate projections, each $[d_{model}, d_{head}]$ .

Concretely: if you have H=32 query heads and G=8 groups, each group serves H/G=4 query heads. The KV cache is G times smaller than MHA — 8× smaller here, versus 32× smaller with MQA.

The attention computation per group

For a group $g$ serving heads $h \in \mathcal{G}_g$ :

\text{Attention}(Q_h, K_g, V_g) = \text{softmax}\!\left(\frac{Q_h K_g^\top}{\sqrt{d_{head}}}\right) V_g

Every head in the group uses the same $K_g$ and $V_g$ , but because $Q_h$ differs per head, the attention pattern and output differ per head. The model learns to spread its representational load across the G groups — a constraint it adapts to during training, which is why GQA benefits most when trained from scratch rather than retrofitted.

Uptraining from MHA

It's possible to adapt a trained MHA checkpoint to GQA by mean-pooling the H KV heads into G groups and then continuing training briefly. The Ainslie et al. paper showed this "uptraining" on roughly 5% of the original training tokens recovered most of the quality gap, making it a practical option for adapting existing models.

Show Me the Code

Measuring the KV cache for MHA, GQA, and MQA at model-serving scale.

1import numpy as np2
3seq, d_model, H, d_head = 4096, 4096, 32, 128  # LLaMA-style4
5def kv_cache_gb(n_groups, seq_len, d_head, n_layers=32, dtype_bytes=2):6    """KV cache in GB for both K and V across all layers."""7    return 2 * n_groups * seq_len * d_head * n_layers * dtype_bytes / 1e98
9mha = kv_cache_gb(H, seq, d_head)10gqa_8 = kv_cache_gb(8, seq, d_head)11gqa_4 = kv_cache_gb(4, seq, d_head)12mqa = kv_cache_gb(1, seq, d_head)13
14print(f"MHA (32 groups):  {mha:.2f} GB")15print(f"GQA  (8 groups):  {gqa_8:.2f} GB  ({mha/gqa_8:.1f}× smaller than MHA)")16print(f"GQA  (4 groups):  {gqa_4:.2f} GB  ({mha/gqa_4:.1f}× smaller than MHA)")17print(f"MQA  (1 group):   {mqa:.2f} GB  ({mha/mqa:.1f}× smaller than MHA)")18# MHA (32 groups):  2.15 GB19# GQA  (8 groups):  0.54 GB  (4.0× smaller than MHA)20# GQA  (4 groups):  0.27 GB  (8.0× smaller than MHA)21# MQA  (1 group):   0.07 GB  (32.0× smaller than MHA)

At 4,096 tokens GQA with 8 groups still needs 0.54 GB per request for the KV cache alone — relevant at high concurrency, where the number of active caches multiplies by the number of simultaneous users.

Watch Out For

Treating group count as a quality knob you can tune at inference time

The number of KV groups is baked in during training. A model trained with 8 groups cannot be re-run with 4 or 16 groups without retraining. Group count is an architectural choice, not a serving parameter. Set it before you train.

Expecting uptraining to fully close the gap

Uptraining an MHA checkpoint to GQA works but leaves a small quality gap that only full training from scratch closes. For a production model where quality matters, train with GQA from the start. Reserve uptraining for research or resource-constrained adaptation.

The Quick Version

GQA divides H query heads into G groups; each group shares one KV pair instead of having one per head (MHA) or sharing one across all heads (MQA).
KV cache is G times smaller than MHA, so G=8 out of H=32 gives a 4× reduction.
Quality is near-MHA quality — the quality drop relative to MHA is measurably smaller than MQA's.
Group count is fixed at training time, not a serving dial.
Most modern open-weight models (Llama 2 70B, Mistral, Gemma 2) use GQA as the default.