Vol. VII · Modern Architectures Plate 2 of 10grouped-query attention, GQA, multi-query attention, MQA, KV cache, key-value headsAC

Fewer Keys and Values.

Every head needs its own question, but they can share the answers’ index. The cache shrinks four- or eight-fold.

In standard multi-head attention every head has its own queries, keys and values. During generation the keys and values of every earlier token are cached for every head (plate VI.9), and that cache, not the weights, is what limits how many sequences a server can hold. Multi-query attention (Shazeer, 2019) keeps a separate query per head but a single shared key and value head: the cache shrinks by the number of heads. Grouped-query attention (Ainslie and colleagues, 2023) sits between: heads are divided into groups, each sharing one key–value head.

Llama 3 8B has 32 query heads and 8 key–value heads, a group size of 4; Mistral 7B has the same ratio. The workshop computes the cache saving for published configurations and shows three recorded runs of micro’s shape with 4, 2 and 1 key–value heads for its 4 query heads.

Below: left, which key–value head each query head reads under the chosen sharing; right, held-out loss of recorded runs of micro’s shape with standard attention (4 key–value heads), grouped (2) and multi-query (1), and the cache each would need per token.

What you will learn

Describe multi-query and grouped-query attention.
Compute the cache saving from the sharing factor.
Measure the quality cost from recorded runs.
Read H and H_KV in published configurations.

Workshop

Left: query heads (top) and the key–value head each reads (bottom) for the chosen number of key–value heads. Right: held-out loss of recorded runs with 4, 2 and 1 key–value heads for 4 query heads.

Shared keys, separate questions

In the recorded runs, halving the key–value heads from 4 to 2 changes the held-out loss from 4.208 to 4.218 nats (0.010), and a single shared head gives 4.230 (0.022). The cache shrinks two- and four-fold, and the attention parameters fall too, from 262,144 to 163,840 in micro. Those are the trade-offs reported at scale: multi-query attention can cost noticeable quality, grouped-query attention with a modest group size costs almost none, and the saving in cache memory converts directly into more sequences per server and lower cost per token.

That is why grouped-query attention became standard almost immediately. Llama 3 8B shares each key–value head among 4 query heads; the GQA paper showed an existing multi-head model could even be converted by averaging its key–value heads into groups and training briefly. Latent attention, the next plate, pushes the same idea further by compressing keys and values into a small shared vector instead of discarding heads.

Multi-query attention was proposed by Noam Shazeer in a 2019 paper whose title — “One Write-Head is All You Need” — echoed the Transformer’s; it was adopted first in PaLM (2022), and grouped-query attention generalised it the following year.

Terms

Multi-query attention (MQA): All query heads share one key–value head.
Grouped-query attention (GQA): Query heads in groups, each group sharing a key–value head.
Key–value head: One set of key and value projections; the unit the cache stores.
Group size: Query heads per key–value head, H / H_KV.

Sharing key–value heads among query heads shrinks the cache by H / H_KV for a small quality cost; grouped-query attention is now standard.

Sources

Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need”, arXiv:1911.02150, 2019.
Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, EMNLP 2023.