Multi-Query Attention (MQA) changes one thing about standard multi-head attention, and that one change buys an enormous amount of inference speed. In ordinary multi-head attention every head has its own query, key, and value projections. MQA keeps all h query heads but forces them to share a single key head and a single value head. The keys and values — the tensors you cache during generation — shrink by a factor of h, and because token-by-token decoding is bottlenecked on reading memory rather than doing arithmetic, a smaller cache translates almost directly into faster tokens and bigger batches. The catch is that collapsing the keys and values to one head throws away representational diversity, which costs a little quality and can make training less stable — the tension Grouped-Query Attention later softened. This piece works through the shapes, derives the cache reduction, runs a concrete example, and explains why MQA was the first, most aggressive step down this path.
The baseline: what multi-head attention actually stores
Start from the standard setup so the change is unambiguous. A model has hidden size d_model and h attention heads, each of dimension d_k = d_model / h. For a token with hidden vector x: [d_model], each head i computes its own query, key, and value:
q_i = x · W_Q^i k_i = x · W_K^i v_i = x · W_V^i
W_Q^i, W_K^i, W_V^i : [d_model, d_k] for i = 1 .. hSo multi-head attention (MHA) produces h distinct query, key, and value vectors per token. That diversity is the point: each head can specialize on a different relationship in the sequence. The cost shows up during generation. To attend to the past without recomputing it, you cache every previous token’s keys and values — all h heads, at every layer. The KV cache grows with context length, and in MHA it grows with the full head count. MQA attacks precisely this term.
The single change: one key head, one value head
MQA keeps the query side untouched — still h independent query projections — but replaces the h key and value projections with exactly one of each, shared across all heads:
q_i = x · W_Q^i i = 1 .. h (h query heads)
k = x · W_K (ONE shared key head)
v = x · W_V (ONE shared value head)
head_i = softmax( q_i · K^T / sqrt(d_k) ) · V (K, V shared for all i)Every query head q_i now attends against the same k and combines the same v; only the query distinguishes one head’s output from another’s. The query-key math per head is identical to MHA; what disappears is the per-head key and value. The parameter savings are modest, but the tensor you must keep around during decoding — the cache of past K and V — is now one head’s worth instead of h. That is where the win lives.
The KV cache, and why decode is memory-bound
To see why shrinking the cache matters so much, recall how autoregressive decoding spends its time. Generating one token is a forward pass that produces a single new position. The arithmetic is tiny — a handful of matrix-vector products — but to do it the hardware must stream the entire set of model weights and the entire KV cache out of memory. The compute units finish long before the bytes arrive.
This is the memory-bandwidth-bound regime: time per token is governed roughly by bytes_read / memory_bandwidth, not by FLOPs. As context grows, the KV cache becomes a large and growing share of those bytes. Prefill (the prompt, processed in parallel) is compute-bound and barely cares about cache size; decode is the opposite. So any optimization that reduces bytes read per decode step pays off directly in tokens per second. Cut the KV cache by a factor of h and you cut the dominant memory-traffic term of long-context generation by that factor — while freeing memory to run larger batches, raising throughput further.
Deriving the cache reduction: a factor of h
Count the bytes. For a model with L layers, hidden size d_model, h heads of size d_k = d_model / h, keys and values at b bytes per element, the per-token cache is:
MHA cache/token = 2 × L × h × d_k × b
= 2 × L × d_model × b (since h · d_k = d_model)
MQA cache/token = 2 × L × 1 × d_k × b
= 2 × L × (d_model / h) × bThe leading 2 counts keys and values separately. Taking the ratio, everything cancels except the head count:
MHA / MQA = (h × d_k) / d_k = hSo MQA reduces the KV cache by exactly a factor of h, the number of heads. Total cache for N tokens and batch B is this per-token figure times N × B, so the same shrink applies to the whole cache. MQA is a pure memory optimization aimed at the one tensor that scales with context.
A worked example
Take a mid-size model: d_model = 4096, h = 32 heads, so d_k = 128; L = 32 layers; keys and values in fp16, b = 2 bytes.
MHA cache/token = 2 × 32 × 4096 × 2 B = 524,288 B ≈ 512 KB
MQA cache/token = 2 × 32 × 128 × 2 B = 16,384 B ≈ 16 KB
reduction = 512 / 16 = 32 = h ✓Now scale to a context of N = 8192 tokens for one sequence:
MHA: 512 KB × 8192 ≈ 4.0 GB per sequence
MQA: 16 KB × 8192 ≈ 128 MB per sequenceFour gigabytes of cache for one 8K conversation is the kind of number that decides whether you can batch at all; 128 MB is almost an afterthought. On memory-bandwidth-bound decode, reading 128 MB instead of 4 GB per step is the difference between a fast responder and one that stalls on memory — and the freed capacity packs roughly h times as many sequences into memory.
Why the speedup is real, not just smaller numbers
Be precise about the mechanism, because MQA does not reduce the number of attention scores computed. Each of the h query heads still forms a full q_i · K^T score vector over the context and still does a weighted sum of values — the attention FLOP count is essentially unchanged.
What changes is arithmetic intensity — FLOPs per byte moved. In decode, the shared K and V are read once and reused by all h query heads, so the same compute is fed by h× fewer cache bytes. When you are waiting on memory, doing the same math against far less traffic is a direct latency win. It compounds with batching: because each sequence’s cache is h× smaller, more sequences fit and the fixed cost of streaming the weights is amortized over more tokens. Smaller cache, higher intensity, bigger batches — three reinforcing reasons decode gets faster.
The cost: quality loss and training instability
Nothing is free. Collapsing h key heads and h value heads down to one removes representational capacity from attention. In full MHA, different heads can key on different features — syntax here, coreference there — because each has its own K and V subspace. In MQA every head must retrieve from the same single key/value subspace; only the query differs. Empirically this produces a measurable, if usually modest, drop in model quality versus an otherwise identical MHA model.
There is a second, subtler problem: training instability. Models trained with MQA, or converted to it, were reported to be more finicky — more sensitive to hyperparameters and prone to quality degradation — than MHA counterparts. The single shared KV head is a bottleneck every query head leans on, and concentrating that much responsibility in one subspace makes the optimization landscape less forgiving. MQA hands you a large, reliable inference win for a real but smaller quality-and-stability tax.
From MQA to GQA: why the extreme came first
MQA is one end of a spectrum. Think of it as a dial on how many key/value heads you keep: MHA keeps all h, MQA keeps exactly 1. Grouped-Query Attention (GQA) is the general case in between — partition the h query heads into g groups and give each group its own shared KV head, so there are g key/value heads with 1 ≤ g ≤ h. MQA is simply GQA at g = 1; MHA is GQA at g = h.
MQA came first (Shazeer, 2019) because it is the most aggressive, simplest point on that dial: one KV head, maximum cache savings, minimum bookkeeping. It proved the thesis — that decode is memory-bound and shrinking the KV cache buys speed. GQA (Ainslie et al., 2023) is the refinement: keeping a handful of KV heads (say g = 8) recovers most of MHA’s quality and stability while still cutting the cache by h / g. MQA was the bold first step that made the trade-off legible; GQA tuned it into the default. The companion GQA article works through that grouped math in full.
Shapes, complexity, and where MQA sits
Collecting the shapes makes the picture compact. Per token, per layer, at head dimension d_k and head count h:
| Quantity | MHA | MQA |
|---|---|---|
| Query heads | h | h |
| Key/value heads | h | 1 |
| KV cache / token / layer | 2 · h · d_k | 2 · d_k |
| Attention FLOPs | O(h · N · d_k) | O(h · N · d_k) |
| Relative cache | 1× | 1/h× |
The rows say it all: query-side compute is untouched, the attention FLOPs are the same order, and the only column that moves is the cache, which drops by h. That is the signature of a memory optimization — it does not make the math cheaper, it makes the math cheaper to feed, the same family as FlashAttention.
Practical implications and pitfalls
For serving small models on constrained hardware — the CPU-SLM case — MQA is especially attractive because CPUs have far less memory bandwidth than GPUs, so the memory-bound decode penalty is sharper and cache reduction matters more; a factor-of-h smaller cache can be what lets a long-context model fit at all. A few cautions. First, the win is a decode/serving win; it does little for compute-bound prefill, so quote your speedups on generation, not prompt processing. Second, you usually cannot bolt MQA onto a pretrained MHA model for free — you either train with it from the start or uptrain (convert and fine-tune), and skipping that step is where quality quietly collapses. Third, if the quality drop is unacceptable, move up the dial to GQA rather than forcing MQA — which is why most modern models chose g > 1. MQA remains the clearest way to see the trade-off: one KV head, cache over h, and a small bill for the privilege.