Chapter 72 — GQA: shrinking the KV cache without losing (much) quality

In one sentence

Grouped-Query Attention keeps many query heads but computes far fewer key/value heads, so groups of query heads share the same K/V pair — shrinking the KV cache (Chapter 38's memory bottleneck) by exactly the sharing ratio, at a small, measured quality cost.

The problem

Chapter 25's multi-head attention gives every head its own W_Q, W_K, W_V — for H heads, that's H separate key and value projections. Chapter 38's KV cache formula counts n_heads for exactly this reason:

KV cache size = 2 x n_layers x n_heads x d_head x seq_len x batch x bytes

n_heads multiplies the cache size just as much for K/V as it does for queries — but empirically, many attention heads end up computing fairly similar attention patterns even though they were free to specialise differently. That redundancy is exactly what GQA exploits.

The spectrum

Scheme	Key/value heads	Cache size	Quality
MHA (Chapter 25's version)	`n_kv_heads = n_heads`	Largest	Best (the baseline)
MQA (Multi-Query Attention)	`n_kv_heads = 1`	Smallest	Noticeably worse — tried first, mostly superseded
GQA (Grouped-Query Attention)	`1 < n_kv_heads < n_heads`	In between	Nearly identical to MHA — the practical default in almost every current LLM

How it actually works

Queries stay at full richness — still H separate W_Q projections. Only the key/value side shrinks to G projections, where G < H. Split the H query heads into G equal groups; every query head in a group attends using that group's single shared K and V:

Normal MHA (H=8 example):
  Q1 K1 V1   Q2 K2 V2   Q3 K3 V3   Q4 K4 V4   Q5 K5 V5   Q6 K6 V6   Q7 K7 V7   Q8 K8 V8
  8 separate K/V pairs -- one per query head

GQA with G=2 groups:
  Q1 Q2 Q3 Q4  -->  share  -->  K_A V_A
  Q5 Q6 Q7 Q8  -->  share  -->  K_B V_B
  Only 2 K/V pairs, computed and cached -- 4x smaller cache

Llama-3-8B's real configuration: 32 query heads, 8 key/value heads — a group size of 4 (every 4 query heads share one K/V pair). This is the actual config used to train and serve that model, not a simplified textbook example.

Worked example: extending Chapter 38's own numbers

Chapter 38 computed the KV cache for a 13B model (40 layers, 40 heads, 128 dims/head, bf16) at plain MHA:

per token, per user:  2 x 40 x 40 x 128 x 2 bytes  =  819,200 bytes  ~= 0.8 MB
4,000-token conversation, 32 users:  ~100 GB

Apply GQA with 8 key/value heads instead of 40 (a 5x reduction in the n_kv_heads term — queries stay at 40 heads, only K/V shrink):

per token, per user:  2 x 40 x 8 x 128 x 2 bytes  =  163,840 bytes  ~= 0.16 MB
4,000-token conversation, 32 users:  ~20 GB          (was ~100 GB)

The same model, the same context length, the same number of users — 80 GB of memory recovered, purely by sharing key/value projections across query heads. This is precisely why Chapter 38 called GQA "why almost every modern LLM uses" it: it directly attacks the term Chapter 38 identified as the usual binding constraint on how many concurrent users a deployment can serve.

Why quality barely drops

Each query head still gets its own learned projection and therefore its own query — what it's looking for doesn't change. Only what it can look at per group becomes shared. In practice, heads within a randomly-assigned group turn out to attend to similar-enough context that little is lost; published ablations (see the GQA paper below) show GQA reaching within a fraction of a point of full MHA quality on standard benchmarks, while MQA's harsher sharing (all heads, one K/V) shows a more visible drop — which is exactly why GQA displaced MQA as the practical default rather than the more extreme option winning outright.

Practice (15 min) — measure the exact cache reduction yourself

def kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, batch, bytes_per_num=2):
    return 2 * n_layers * n_kv_heads * d_head * seq_len * batch * bytes_per_num

configs = [
    ("Plain MHA (40 KV heads)", 40, 40),
    ("GQA (8 KV heads)",        40, 8),
    ("GQA (4 KV heads)",        40, 4),
    ("MQA (1 KV head)",         40, 1),
]
for name, n_layers, n_kv_heads in configs:
    size = kv_cache_bytes(n_layers, n_kv_heads, d_head=128, seq_len=4000, batch=32)
    print(f"{name:28s} {size/1e9:6.2f} GB")

Run it and confirm the 5x and 40x reductions land exactly where the maths above predicts — this is a case where the "engineering trick" is genuinely just arithmetic, not a black box.

Common confusions

GQA is not a form of quantization (Chapter 73) — it changes the architecture (how many K/V projections exist), not the precision of the numbers stored. They compose: a GQA model's already-smaller cache can be quantized further (Chapter 38's "KV cache quantisation" row). It's also not distillation (Chapter 67) — GQA is a design decision made during pretraining, not a post-hoc compression of an existing model.
Queries are never shared, only K/V. It's tempting to assume GQA "reduces the number of heads" generally — it doesn't touch the query side at all, which is why quality holds up as well as it does: each head still asks its own question, it just shares a smaller pool of things to look at.
The group size is a training-time architectural choice, fixed when the model is pretrained — you cannot bolt GQA onto an already-trained plain-MHA model without retraining or a dedicated "uptraining" procedure.

Check yourself

In GQA, which projections shrink (Q, K, V, or some combination) and which stay full-sized?
Using Chapter 38's 13B-model numbers, what's the KV cache size at 4 concurrent users with 8 KV heads instead of 40?
Why does GQA lose less quality than MQA?
Why can't you add GQA to a model after it's already been pretrained with plain MHA, the way you might apply quantization after training?