Grouped-Query Attention: one set of keys for several heads

Ainslie et al., 2023 — and Shazeer, 2019

The whole implementation is one repeat. Getting it in the wrong place produces plausible garbage and no error.

Download gqa_from_scratch.py — everything below runs from it. NumPy only, under a second.

I read the KV-cache page — I wrote it — and it ends on arithmetic. The cache is what stops you serving the model; shrink kv_heads and the cache shrinks by exactly the same factor. Fine. Then I opened a real implementation to see what "shrink kv_heads" is as code, and it was six lines, four of them a tensor reshape I couldn't read. Nobody told me the change is one repeat in the right place, and that the whole class of bugs here is putting it in the wrong one.

A read replica shared by several app servers

You've sized this before:

app_server_0 -> replica_0     # one replica each. Perfect isolation,
app_server_1 -> replica_1     # and you pay for eight replicas.
...
app_server_7 -> replica_7

Replicas are the expensive part, so you consolidate. Eight app servers, two replicas, four servers per replica: 

app_server_0..3 -> replica_0   # each server still does its own work,
app_server_4..7 -> replica_1   # but they read the same snapshot

That is grouped-query attention, exactly. Query heads are the app servers — every one keeps its own Wq, its own scores, its own output. Key and value heads are the replicas, and they're what you store. Multi-head is one replica per server. Multi-query is one replica for the whole fleet. GQA is the middle setting, and like every capacity decision you've made, the question is how much shared staleness you can tolerate.

TermMeansProgrammer's analogy
Query headOne independent "what am I looking for". Never cached — used once, thrown awayAn app server. Cheap to add
KV headOne set of keys and values, cached for every token so farA read replica. The thing you pay for
GroupThe n_heads / n_kv_heads query heads sharing one KV headThe servers pinned to one replica
MQA (Shazeer, 2019)n_kv_heads = 1One replica, whole fleet
GQA (Ainslie et al., 2023)1 < n_kv_heads < n_headsThe middle setting

Why anyone did this

Not for FLOPs. Decoding is memory-bound: each step drags the weights and the whole cache across the memory bus to do a rounding error's worth of arithmetic. The cache is the term that grows with every token and every concurrent user, and kv_heads is the only factor in it that the architecture gets to choose. The KV-cache page has that table — bytes per token, the 4× and 32× reductions, the point where the cache outgrows the checkpoint. I'm not reprinting it. This page is the forward pass and what the sharing costs.

Multi-query came first: Shazeer, 2019, "Fast Transformer Decoding: One Write-Head is All You Need", a short paper whose whole argument is decoding speed. GQA arrived in 2023 with two claims — that the useful setting is in between, and that you don't have to pre-train it. You can convert a checkpoint you already have.

The one formula

headi = Attention(Qi, K⌊i/g⌋, V⌊i/g⌋),   g = nheads / nkv

That's the entire paper as notation. The subscript on K and V is the only thing that changed since 2017, where it read Ki, Vi. Read the layer as pipeline stages:

StepShapeWhat it does
X @ W_q(n_heads, n, head_dim)All the query heads. Unchanged in every variant
X @ W_k, X @ W_v(n_kv, n, head_dim)The narrower projections. This is the change, and this is what gets cached
repeat_kv_heads→ (n_heads, n, head_dim)The whole implementation. Expand each KV head g times so the queries have something to meet
Q @ KT / √d(n_heads, n, n)Ordinary scaled dot-product. Full head count, full cost
softmax @ V(n_heads, n, head_dim)Weighted blend, per head
merge @ W_o(n, d_model)Concatenate heads, project out. Also unchanged
Note where the repeat sits: after the projection, before the matmul, and never on the cache. You store the narrow K and V and expand them at compute time. Expand first and you've written multi-head attention with extra steps.

One function, three variants

The three names are one function with one argument. Pass n_kv_heads = n_heads and it's multi-head. Pass 1 and it's multi-query. Pass anything in between and it's grouped-query. The script prints the shape at every step so you can see that the only line that differs is the K/V projection width — everything downstream is identical.

Don't take my word for it. Here's the actual output:

  n_kv_heads=8  multi-head  (MHA, 2017)
    X (6, 64) -> Q (8, 6, 8)  K (8, 6, 8)  V (8, 6, 8)
    Q (8, 6, 8)  K after repeat (8, 6, 8)  V after repeat (8, 6, 8)
    scores (8, 6, 6)  weights (8, 6, 6)  out (8, 6, 8)
    layer output (6, 64)   <- identical for all three

  n_kv_heads=2  grouped-query (GQA, 2023)
    X (6, 64) -> Q (8, 6, 8)  K (2, 6, 8)  V (2, 6, 8)
    Q (8, 6, 8)  K after repeat (8, 6, 8)  V after repeat (8, 6, 8)
    scores (8, 6, 6)  weights (8, 6, 6)  out (8, 6, 8)
    layer output (6, 64)   <- identical for all three

  n_kv_heads=1  multi-query (MQA, 2019)
    X (6, 64) -> Q (8, 6, 8)  K (1, 6, 8)  V (1, 6, 8)
    Q (8, 6, 8)  K after repeat (8, 6, 8)  V after repeat (8, 6, 8)
    scores (8, 6, 6)  weights (8, 6, 6)  out (8, 6, 8)
    layer output (6, 64)   <- identical for all three

One line differs across the three blocks: the K and V shapes before the repeat. That's the paper. Everything after the repeat is the 2017 code you already know.

The repeat, and the bug nobody catches

You have 2 KV heads and need 8. There are two obvious ways to stretch a list of 2 into a list of 8, and NumPy gives you both: 

np.repeat(kv, 4, axis=0)   -> [kv0 kv0 kv0 kv0 kv1 kv1 kv1 kv1]  # block
np.tile(kv, (4, 1, 1))     -> [kv0 kv1 kv0 kv1 kv0 kv1 kv0 kv1]  # interleaved

Same shape. Same dtype. No error, either way. But every checkpoint on disk was saved under the convention that query head i belongs to KV head i // g — contiguous blocks — so only the first one is right. Llama's reference implementation calls the operation repeat_kv and documents it as torch.repeat_interleave over the head axis, which is the block version. 

  n_heads=8, n_kv_heads=2, group=4
  query head       : [0, 1, 2, 3, 4, 5, 6, 7]
  block  -> kv head: [0, 0, 0, 0, 1, 1, 1, 1]  matches i // group
  interl -> kv head: [0, 1, 0, 1, 0, 1, 0, 1]  matches i %  n_kv
  convention (i // group): [0, 0, 0, 0, 1, 1, 1, 1]

  output shapes equal?            True
  no exception raised?            True
  max |block - interleave|        0.2661
  mean |block|                    0.0651
  -> same shape, no error, different answer. That is the bug.
The error is four times the size of the signal. Not a crash, not a NaN, not a shape mismatch — a model that loads, runs, and talks nonsense. If you ever port weights between frameworks and the outputs are "close but wrong", this is the first thing to check.

At n_kv = n_heads, it is exactly multi-head

Worth proving rather than asserting, because it's how you sanity-check any GQA implementation you write: set the group size to 1 and the repeat becomes a no-op, so the function must reduce to plain 2017 attention. The script computes multi-head the long way, with no repeat call anywhere, and diffs. 

  max |grouped(n_kv=8) - plain multi-head| = 0.000e+00
  bit-identical?                            True
  group size is 1, so np.repeat is a no-op. Same code path.

Not "to float precision" — bit-identical, because np.repeat with a count of 1 returns the same numbers in the same order. Your first test for any implementation of this: does it still match multi-head when you turn the sharing off?

What sharing actually costs

Now the honest part. Heads exist because they look for different things. Force four of them through one set of keys and they can't all get what they want. To show that mechanically I build eight query heads whose ideal key is a different token each, then mean-pool the key heads into fewer and fewer groups and measure how far the output drifts from the ideal. 

    n_kv | group |  max |out - ideal| | mean |out - ideal|
       8 |     1 |           0.0000 |           0.0000
       4 |     2 |           0.0321 |           0.0092
       2 |     4 |           0.5660 |           0.1723
       1 |     8 |           1.4787 |           0.4548

  per-head error at n_kv=2 (heads 0-3 share, heads 4-7 share):
    head 0  group 0  error 0.3688   argmax attends to token 0 (wanted 0)
    head 1  group 0  error 0.4308   argmax attends to token 1 (wanted 1)
    head 2  group 0  error 0.3620   argmax attends to token 2 (wanted 2)
    head 3  group 0  error 0.4960   argmax attends to token 3 (wanted 3)
    head 4  group 1  error 0.2810   argmax attends to token 4 (wanted 4)
    head 5  group 1  error 0.3821   argmax attends to token 5 (wanted 5)
    head 6  group 1  error 0.2884   argmax attends to token 6 (wanted 6)
    head 7  group 1  error 0.5660   argmax attends to token 7 (wanted 7)

The shape of it is what matters. Halving the KV heads is nearly free (0.0092); quartering costs an order of magnitude more; collapsing to one costs another factor of three. The heads still point at the right token — the argmax is right in every row — but the distribution around it has been dragged toward the group average, and the output carries that. That's the graceful-degradation curve GQA is trading on.

This is a mechanical illustration on vectors I constructed to make the point, not a quality measurement. The paper's quality claims — that GQA lands close to multi-head while MQA loses more — are the paper's, measured on real tasks with real uptrained checkpoints. Don't quote my numbers as evidence about language models; they're evidence about averaging.

Parameters shrink a little. FLOPs don't move at all.

The interview question here is "so GQA is faster because it does less arithmetic, right?" — and no. The keys are expanded back to the full head count before the matmul, so the attention does exactly the same work. Wq and Wo are untouched. Only Wk and Wv get narrower, which is a rounding error against total parameters — and the cache, which is not. 

  Llama-2-7B-shaped layer: d_model=4096, n_heads=32, head_dim=128, seq=2048

    variant     | n_kv | W_k+W_v params | vs MHA | attn MACs/token
    multi-head  |   32 |     33,554,432 |   1.0x | 16,777,216
    GQA-8       |    8 |      8,388,608 |   4.0x | 16,777,216
    GQA-4       |    4 |      4,194,304 |   8.0x | 16,777,216
    multi-query |    1 |      1,048,576 |  32.0x | 16,777,216

  W_q and W_o are untouched: 16,777,216 params each, in every variant.

Same column, four times. The win is memory traffic, not maths: fewer bytes to read per decode step, and a cache small enough to batch more users against. Bytes per token, and how fast that becomes your whole hardware bill, is the KV-cache page.

Uptraining: the part of the paper everyone forgets

The paper is titled Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, and that title is the contribution. Nobody wants to pre-train a 70B model again to change one hyperparameter. So: take the multi-head checkpoint you already have, mean-pool the key and value heads inside each group, and train briefly — the paper reports doing it with a small fraction of the original pre-training compute.

Mean-pooling, specifically, and not "keep head 0". W_k is (d_model, n_heads × head_dim); view it as (d_model, n_kv, group, head_dim), average over the group axis, flatten. Four lines. The mean is the vector closest to all group members at once; selecting one throws the other g−1 away. So mean-pooling starts the short training run from a better place. That's the whole argument, and it's measurable: 

  original multi-head output, mean |y| = 0.0673

    n_kv | mean-pooled err | first-head err | first/mean
       4 |          0.0414 |         0.0671 |       1.62x
       2 |          0.0580 |         0.1005 |       1.73x
       1 |          0.0653 |         0.0794 |       1.22x

Mean-pooling is closer at every group size — up to 1.7× closer. Note also that at n_kv=1 the error is the same size as the signal itself: conversion alone does not give you a working model. It gives you an initialisation. The brief training run is not optional, it's the point. That distinction is the thing to have straight if someone asks you how Llama-2 70B got its GQA.

Where this landed. Llama-2's 70B model uses GQA, and most open models released since have followed. The notable dissent is DeepSeek, which compresses K and V into a shared latent vector instead of sharing heads — that's multi-head latent attention, a different answer to the same cache problem. The rest of the Llama recipe — RMSNorm, SwiGLU, RoPE — is on its own page, and the attention this all modifies is here.

The code

Pure NumPy. No torch, no autograd, no training — the forward pass and the checkpoint conversion, so you can watch the numbers move. Every function under 20 lines; the whole thing runs in well under a second. Everything below is verbatim from the script, minus softmax and the two reshape helpers.

Download the full script — about 150 lines of mechanism plus seven runnable demos. Then: 

python3 gqa_from_scratch.py
# THE WHOLE IMPLEMENTATION. One repeat, in the right place.
#
# kv is (n_kv_heads, n, head_dim). We need (n_heads, n, head_dim) so
# every query head has a key head to talk to.
#
#   "block"      np.repeat -> [kv0 kv0 kv0 kv0 kv1 kv1 kv1 kv1]
#                             head i uses kv head i // group  <-- right
#   "interleave" np.tile   -> [kv0 kv1 kv0 kv1 kv0 kv1 kv0 kv1]
#                             head i uses kv head i %  n_kv   <-- wrong
#
# Both produce the correct SHAPE. Only one matches the convention every
# checkpoint on disk was saved with. Getting it wrong raises nothing.
def repeat_kv_heads(kv, n_heads, mode="block"):
    n_kv_heads = kv.shape[0]
    if n_heads % n_kv_heads != 0:
        raise ValueError(
            f"n_heads={n_heads} not divisible by n_kv_heads={n_kv_heads}: "
            "groups must be equal-sized"
        )
    group = n_heads // n_kv_heads
    if mode == "block":
        return np.repeat(kv, group, axis=0)   # head i -> kv head i // group
    if mode == "interleave":
        return np.tile(kv, (group, 1, 1))     # head i -> kv head i % n_kv
    raise ValueError(f"unknown repeat mode {mode!r}")


# This ONE function is multi-head, multi-query and grouped-query. The
# only difference is how many rows K and V arrive with.
#
#   Q  (n_heads,    n, head_dim)
#   K  (n_kv_heads, n, head_dim)   n_kv == n_heads -> multi-head
#   V  (n_kv_heads, n, head_dim)   n_kv == 1       -> multi-query
#                                  in between      -> grouped-query
def attention(Q, K, V, mode="block", causal=False, verbose=False):
    n_heads, n, head_dim = Q.shape
    K = repeat_kv_heads(K, n_heads, mode)
    V = repeat_kv_heads(V, n_heads, mode)

    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(head_dim)   # (H, n, n)
    if causal:
        mask = np.tril(np.ones((n, n), dtype=bool))
        scores = np.where(mask, scores, -np.inf)
    weights = softmax(scores, axis=-1)
    out = weights @ V                                       # (H, n, dim)
    return out, weights


# One attention layer. n_kv_heads picks the variant. W_q and W_o do not
# move; only W_k and W_v get narrower.
def forward(X, W, n_heads, n_kv_heads, head_dim, mode="block"):
    Q = split_heads(X @ W["W_q"], n_heads, head_dim)
    K = split_heads(X @ W["W_k"], n_kv_heads, head_dim)
    V = split_heads(X @ W["W_v"], n_kv_heads, head_dim)
    out, _ = attention(Q, K, V, mode=mode)
    merged = out.transpose(1, 0, 2).reshape(X.shape[0], n_heads * head_dim)
    return merged @ W["W_o"]


# UPTRAINING, the paper's actual contribution. You do not train a GQA
# model from scratch. You take an existing multi-head checkpoint and
# MEAN-POOL the key and value heads inside each group, then train
# briefly. The mean is the projection closest to all group members;
# picking one head throws the other g-1 away.
def convert_checkpoint(W_mha, n_heads, n_kv_heads, head_dim, how="mean"):
    group = n_heads // n_kv_heads
    out = {"W_q": W_mha["W_q"].copy(), "W_o": W_mha["W_o"].copy()}
    for name in ("W_k", "W_v"):
        w = W_mha[name].reshape(-1, n_kv_heads, group, head_dim)
        if how == "mean":
            w = w.mean(axis=2)            # mean-pool inside each group
        elif how == "first":
            w = w[:, :, 0, :]             # keep head 0, bin the rest
        out[name] = w.reshape(-1, n_kv_heads * head_dim)
    return out

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — reverse-engineering from broken behaviour is your native mode, and every row here is a bug someone has actually shipped.

Break thisWatch for
Pass mode="interleave" to forwardNo error, no shape mismatch. max |block − interleave| = 0.2661 against a mean output magnitude of 0.0651 — four times the signal. A model that loads and talks nonsense.
Call repeat_kv_heads with n_kv_heads=3, n_heads=8ValueError: n_heads=8 not divisible by n_kv_heads=3. Uneven groups mean head→group isn't a function; every implementation requires the divide.
In convert_checkpoint, use how="first"Conversion error goes from 0.0580 to 0.1005 at n_kv=2 — 1.73× further from the original model before uptraining even starts.
Share W_k but leave W_v at full widthCache per token per layer goes 2,048 → 5,120 floats instead of 8,192. Your 4× reduction is now 1.6×, for the same quality cost.
Move repeat_kv_heads to before the cache appendCached floats per token per layer back to 8,192 — identical maths, zero saving. You've written multi-head with extra steps. See the cache page.
Set n_kv_heads=1 in demo 4Mean error jumps to 0.4548, roughly 50× the n_kv=4 error. That gap is why GQA exists instead of everyone just using MQA.
Then read the paper — it's short. The method section (the uptraining procedure and the mean-pooling figure) is the part worth your time; the rest is ablations on summarisation and translation benchmarks that you can skim for the shape of the curve rather than the numbers. Read Shazeer's 2019 multi-query paper first if you want the original argument, and note it's mostly about decoder memory bandwidth, not quality.

On the maths anxiety

There isn't any here, and that's not me being encouraging. GQA needs one idea: the mean of g vectors is the single vector closest to all of them. You already believe that. Everything else on this page is array indexing — knowing that head i lives at columns i·head_dim : (i+1)·head_dim, and that i // g maps a head to its group.

What is genuinely hard is the reshape. (n, h·d) → (h, n, d) and back, with a transpose in the middle, is where I lose an afternoon every time. Print .shape after every line — that's how the people who wrote these implementations debug them too. And the judgement call, 8 groups or 4, GQA or latent attention, isn't maths either. It's capacity planning: cache bytes per token against acceptable quality loss, at your context length and batch size.

Where to go next

This page gave you one thing: the forward pass and the conversion, in code you can break. It deliberately skipped the memory arithmetic, the real benchmarks, and everything about how this interacts with a serving stack.

Order I'd take them in: Shazeer first for the motivation, then the GQA paper's method section, then llama/model.py to see it shipped — and if you want the road not taken, latent attention, which compresses instead of shares.