Multi-Query and Grouped-Query Attention

Standard multi-head attention gives every head its own query, key, and value projections. At inference, that means every head needs its own cached key and value vectors for every past token (see kv-cache), and that cache, not the model's weights, is often what actually runs a GPU out of memory during long-context serving. Shazeer, 2019 proposed a specific, surgical fix: keep multiple query heads, but share a single key and value head across all of them.

The asymmetric cut

Multi-query attention (MQA) keeps h separate query heads but collapses the key and value projections down to one shared head each:

# standard multi-head: h separate K, V pairs
K_i, V_i for i in 1..h

# multi-query: one shared K, V for every query head
K, V   (shared)
Q_i for i in 1..h

Every query head still computes its own attention pattern against the same shared keys, so the model retains its ability to run multiple independent matching criteria per position (see multi-head-attention-deep). What it loses is the ability to have differently-projected keys and values per head. The cache that must be stored per token shrinks by a factor of h for keys and h for values, because there is only one key vector and one value vector per token instead of h of each. For long sequences and large batch sizes, where the KV cache dominates memory far more than the model weights do, this is the single largest lever available for fitting more context or more concurrent requests into the same GPU memory.

Why keys and values, not queries

The asymmetry is deliberate. Queries are used once per generated token and then discarded; they never need to be cached, because a new query is computed fresh at every decoding step from the current token only. Keys and values for every past token, in contrast, must be kept around for the entire generation, because every future token's attention needs to compare against all of them. Cutting the query heads down would save nothing at inference (queries were never the memory cost); cutting keys and values down attacks the actual bottleneck directly.

Grouped-query attention: the compromise that shipped

MQA's aggressive sharing (one key/value head for the entire model) sometimes cost noticeable quality relative to full multi-head attention, since every query head is now forced to share the exact same keys and values, with no room for even a few independently projected views. Ainslie et al., 2023 proposed grouped-query attention (GQA): partition the h query heads into g groups, and give each group its own shared key/value head, with g somewhere between 1 (MQA) and h (full multi-head).

g groups, each of size h/g query heads
each group shares one K, V pair
g = 1        -> multi-query attention
g = h        -> standard multi-head attention
1 < g < h    -> grouped-query attention

Ainslie et al. also showed a practical migration path: an existing multi-head model can be converted to GQA by mean-pooling its existing key/value heads within each group and then briefly fine-tuning, rather than retraining from scratch. This is why GQA spread quickly through production models (Llama 2 70B, Llama 3, Mistral, and others use it): labs could retrofit the memory savings onto architectures they had already invested in training.

When it falls down

• It is a quality/memory tradeoff, not a free win. Sharing keys and values across a group necessarily reduces the diversity of matching criteria available to that group's query heads relative to full multi-head attention; the choice of g is a real hyperparameter search, not a default that is always correct.
• MQA's aggressive sharing shows up most on tasks needing precise per-head specialisation. Tasks that lean heavily on fine-grained positional or syntactic distinctions can be more sensitive to collapsing keys and values to a single shared view than average benchmark numbers suggest.
• It only helps the KV cache, not compute. GQA and MQA reduce memory bandwidth and footprint for cached keys and values; they do not reduce the number of floating-point operations in the attention score computation itself, so they are a memory-bound optimisation, not a compute-bound one.
• Retrofitting is lossy. Mean-pooling existing key/value heads and fine-tuning recovers most, not all, of the original model's quality; a model trained with GQA from scratch and a multi-head model retrofitted to GQA are not guaranteed to land at the same point.

Further reading

• Shazeer, 2019, Fast Transformer Decoding: One Write-Head is All You Need, arXiv:1911.02150 - the original multi-query attention proposal.
• Ainslie et al., 2023, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, arXiv:2305.13245 - grouped-query attention and the uptraining recipe.