LM10-T3 — Multi-Query and Grouped-Query Attention

Learning objectives

After this lesson, you should be able to:

• Describe MHA, GQA and MQA by the number of key-value heads.
• Compute attention parameters and KV-cache bytes per token for each.
• Implement grouped-query attention by repeating key-value heads and test it.
• Choose a number of key-value groups for a memory budget and explain the quality risk.

The idea: queries need diversity, keys and values less so

In multi-head attention (Module 3), a layer with H heads projects every token to H queries, H keys and H values. Head j compares its query with the keys of head j and mixes the values of head j. The heads learn different patterns: one follows syntax, one looks at the previous token, one retrieves a name from far back.

During training everything is computed at once, and the number of key and value heads is a question of parameters and arithmetic, nothing dramatic. During generation it is different. At every step the model needs the keys and values of all earlier tokens, and recomputing them would be wasteful, so they are stored: the KV cache (Lesson LM10-T5). The cache has H key heads and H value heads per layer per token and has to be read from memory at every step. It is the largest memory object in long-context serving.

The observation behind MQA and GQA is that the queries carry most of the head diversity. Several query heads can ask different questions of the same set of keys and values. So share the keys and values:

• MHA: H_kv = H. Maximal flexibility, maximal cache.
• MQA (Shazeer, 2019): H_kv = 1. Minimal cache, noticeably lower quality and sometimes unstable training.
• GQA (Ainslie et al., 2023): 1 < H_kv < H, typically 8. Intermediate cache, quality close to MHA in the published experiments.

What does not change

The number of query heads, the width of the layer’s output, the causal mask and the softmax are all unchanged. A group of H/H_kv query heads simply reads the same stored keys and values. The attention computation is the same; only the keys and values are repeated, conceptually, across the group.

Why it also speeds up decoding

A decode step reads the weights and the whole cache from memory (Lesson LM10-T5). A smaller cache means less to read for long contexts, and it allows larger batches. The speed-up is therefore largest for long contexts and many concurrent requests, which is where serving is expensive.

A trade-off, not a free lunch

The mathematics: shapes, parameters and cache

Let d be the model width, H the number of query heads, H_kv the number of key-value heads (dividing H), d_h the head width, and g = H/H_kv the group size. For input X of shape (B, T, d):

Query head j (0-based) uses key-value head ⌊j/g⌋. Equivalently, repeat each key-value head g times along the head dimension (repeat_interleave) to get shape (B, H, T, d_h) and apply ordinary attention:

W_Q

d × (H·d_h); unchanged by sharing.

W_K, W_V

each d × (H_kv·d_h): this is where the parameters are saved.

W_O

(H·d_h) × d; unchanged.

g

group size; query heads j = ig … ig + g − 1 share key-value head i.

Parameters

Cache per token

The 7B configuration (d = 4,096, H = 32, d_h = 128, L = 32, 16-bit):

The cache falls in proportion to H_kv, but the attention parameters fall only to half at the extreme, because the query and output projections (2d² of the original 4d²) are untouched. Sharing is therefore worth it for the cache, not for the parameter count. With a 32-layer model, 8,192-token contexts and 16 concurrent requests, the cache is 68.7 GB with MHA and 17.2 GB with GQA-8 (Lab LM10-L1).

Converting a multi-head checkpoint

A model trained with MHA can be converted: for each group of g heads, average their key projection matrices into one and likewise the value projections (mean pooling), then continue training for a small fraction of the original compute so that the model adapts. The GQA paper reports that this “uptraining” recovers most of the quality at a few percent of the original training cost, and that mean pooling beats keeping only one head of each group. Converting without any further training generally damages quality.

A worked example: implement and test GQA

The first block counts parameters and cache for the layouts of the table. The second implements grouped-query attention for the small shapes of Notebook LM10-L3, and tests three facts: with two key-value heads and eight query heads the output equals ordinary multi-head attention on the repeated heads; query heads of one group do not depend on another group’s keys; and the cache holds a quarter of the numbers.

import torch

def counts(d, nq, nkv, hd, layers, bytes_=2):
    params = d * nq * hd + 2 * d * nkv * hd + nq * hd * d
    cache_token = 2 * layers * nkv * hd * bytes_
    return params, cache_token

for name, nkv in (('MHA', 32), ('GQA-8', 8), ('GQA-4', 4), ('MQA', 1)):
    p, c = counts(4096, 32, nkv, 128, 32)
    print(name, p, c, round(c / 524288, 4))
# MHA 67108864 524288 1.0
# GQA-8 41943040 131072 0.25
# GQA-4 37748736 65536 0.125
# MQA 34603008 16384 0.0312

def attend(q, k, v):
    """q: (B, Hq, T, hd); k, v: (B, Hkv, T, hd). Each key-value head serves Hq // Hkv query heads."""
    rep = q.shape[1] // k.shape[1]
    k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
    scores = q @ k.transpose(2, 3) / q.shape[-1] ** 0.5
    mask = torch.triu(torch.ones(q.shape[2], k.shape[2], dtype=torch.bool), diagonal=1)
    return torch.softmax(scores.masked_fill(mask, float('-inf')), dim=-1) @ v

torch.manual_seed(0)
B, T, hd = 1, 5, 4
q = torch.randn(B, 8, T, hd)
k, v = torch.randn(B, 2, T, hd), torch.randn(B, 2, T, hd)
out = attend(q, k, v)
print(out.shape)                                                # torch.Size([1, 8, 5, 4])

k8, v8 = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
print(torch.allclose(out, attend(q, k8, v8)))                   # True: same as MHA with copied heads

k2 = k.clone(); k2[:, 1] = torch.randn(B, T, hd)                # change the keys of group 1 only
out2 = attend(q, k2, v)
print(torch.allclose(out[:, :4], out2[:, :4]), torch.allclose(out[:, 4:], out2[:, 4:]))   # True False
print(k.numel() + v.numel(), k8.numel() + v8.numel())           # 80 320: a quarter of the stored numbers

Query heads 0 to 3 are untouched when the keys of group 1 change, and heads 4 to 7 change: each query head depends on exactly one key-value head. The cache stores 80 numbers instead of 320 for this tiny example.

Choosing the number of groups

1. Fix the target,Memory left after the weights, the context length and the number of concurrent requests you must serve.
2. Compute the cache,2 · L · H_kv · d_h · b · n · B for candidate values of H_kv (powers of two dividing H).
3. Check the quality,Compare your evaluation scores for the candidates, trained or converted, with the confidence intervals of Module 9.
4. Check the speed,Measure decode throughput at the real batch size and context length; the gain grows with both.

Engineering implications and limitations

Eight is a common choice. Many open models use around 8 key-value heads regardless of size. It cuts the cache by 4 to 8× relative to MHA for 32 to 64 query heads, and the group size is large enough for most of the saving and small enough to keep quality.

Keep the cache before repeating. Store the H_kv heads, not the repeated H. Repeating is a view or a cheap copy that you do at attention time; storing the repeated tensors throws away the whole saving.

Tensor parallelism interacts. When the layer is split across devices by heads, the number of key-value heads must be compatible with the number of devices; with too few heads some devices hold copies of the same one.

Test equality with MHA. Setting H_kv = H must reproduce ordinary multi-head attention exactly. Keep that test in your suite.

Read the configuration. A checkpoint states H_kv in its configuration (often called the number of key-value heads). Loading it into code that assumes MHA gives shape errors at best and wrong text at worst.

Account for both budgets. The parameter saving (about 12% for GQA-8 in the example model) pays for part of a bigger model at the same size; the cache saving pays for concurrency and context length. Report both.

Common misconceptions

Engineering takeaways

1. MHA, GQA and MQA differ only in the number of key-value heads H_kv; the number of query heads stays H.
2. Attention parameters: d·H·d_h + 2·d·H_kv·d_h + H·d_h·d. Cache per token: 2·L·H_kv·d_h·b.
3. In the 7B example GQA-8 keeps 62.5% of the attention parameters and 25% of the cache; MQA keeps 51.6% and 3.1%.
4. Implement by repeating each key-value head g times at attention time; store only H_kv heads in the cache.
5. Convert MHA checkpoints by mean-pooling the heads of each group and uptraining briefly; choose H_kv by memory budget and measured quality.

References and further reading

The parameter and cache numbers in the table were computed by the code in this lesson for a generic 7-billion-parameter configuration. The cited papers give the quality and speed measurements; this lesson does not reproduce them.

• Shazeer, N. (2019). Fast transformer decoding: One write-head is all you need. arXiv:1911.02150.
• Ainslie, J. et al. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. arXiv:2305.13245.
• Touvron, H. et al. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv:2307.09288.
• Ainslie et al. (2023), GQA (research paper)
• Shazeer (2019), Multi-Query Attention (research paper)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: SwiGLU, RMSNorm and mixture of experts, the changes to the feed-forward half of the block.

Related content

Nothing below is required to finish this lesson.

• Before
  LM10-T2: Rotary Position Embeddings

  Rotate query and key pairs by position-dependent angles so attention depends on relative offset.

• Lab
  LM10-L2: Compare Attention Variants

  Count parameters, cache size and compute for multi-head, multi-query and grouped-query attention.