Multi-query attention

Formulation#

Multi-head attention with $h$ heads computes, for head $i$:

$$ O_i = \operatorname{softmax}!\left(\frac{Q_i K_i^{\top}}{\sqrt{d_h}}\right) V_i, \qquad Q_i = X W^Q_i,\; K_i = X W^K_i,\; V_i = X W^V_i $$

$X \in \mathbb{R}^{N \times d}$ is the layer input, $d_h$ the head dimension, and $W^Q_i, W^K_i, W^V_i \in \mathbb{R}^{d \times d_h}$ the per-head projections.

Multi-query attention (Shazeer, 2019, Fast Transformer Decoding: One Write-Head is All You Need, arXiv:1911.02150) keeps $h$ distinct $W^Q_i$ but collapses the others to a single shared pair $W^K, W^V \in \mathbb{R}^{d \times d_h}$:

$$ O_i = \operatorname{softmax}!\left(\frac{Q_i K^{\top}}{\sqrt{d_h}}\right) V, \qquad K = X W^K,\; V = X W^V $$

The output projection is unchanged. Only the number of distinct key/value tensors changes, from $h$ to $1$.

Why this is a bandwidth argument, not a FLOP argument#

Per decoding step, attention FLOPs are $\Theta(h \, d_h \, S)$ under both schemes, because every query head still attends over the whole context. What changes is bytes read:

$$ \text{KV bytes} = 2 \, b \, L \, n_{kv} \, d_h \, S \, B $$

with $b$ bytes per element, $L$ layers, $n_{kv}$ key/value heads, $S$ context length and $B$ batch. MQA sets $n_{kv} = 1$.

Shazeer frames the resulting quantity as memory-access-to-arithmetic ratio. For incremental decoding, standard MHA has ratio $\Theta(n/d + 1/b)$ where $n$ is sequence length and $b$ batch; MQA reduces the dominant term by a factor of $h$. The paper's own summary of the result is that models "can indeed be much faster to decode, and incur only minor quality degradation from the baseline" — read the tables in the paper for the per-configuration figures rather than trusting a single headline speed-up.

The quality cost, honestly#

Shazeer reports a small but real quality loss against MHA on translation. Ainslie et al. (2023) (see grouped-query attention) confirm the gap on summarisation and question answering, and additionally report training instability at scale under MQA — loss spikes that do not occur with intermediate group counts.

The mechanistic story is that heads specialise. Different heads in a trained transformer attend to syntactic dependencies, positional offsets, repeated entities and induction patterns. Forcing one key space to serve all of them removes a degree of freedom the model demonstrably uses.

Deployment record#

Falcon is the clearest illustration. The 7B model shipped with one key/value head; the 40B model in the same release already used eight. The larger the model, the less tolerable MQA's quality cost became.

MQA was a two-year architecture. Its historical importance is that it proved the KV cache, not attention FLOPs, is the binding constraint on decoding, and it created the axis along which GQA and multi-head latent attention were later placed.

Interaction with other techniques#

MQA composes with FlashAttention: one reduces cache size, the other reduces intermediate traffic, and they are orthogonal.

It composes badly with tensor parallelism. Splitting $h$ query heads across $t$ devices leaves a single key/value head that must be replicated on every device, so the per-device cache saving is $h/t$ rather than $h$. GQA with $n_{kv} = t$ is the natural fit, which is one practical reason 8 became the common group count.

Cache quantisation multiplies with all of it. An fp8 KV cache halves $b$ and is close to free in quality terms at 8 bits, though int4 caches degrade measurably on long-context retrieval.

What to learn next#

• Grouped-query attention — the compromise that actually shipped.
• Multi-head latent attention — compressing the cache instead of sharing it.
• Attention — the heads and what they each learn to look for.