Core ML

Multi-Query Attention

All query heads share a single set of keys and values instead of each keeping their own, which shrinks the KV cache by as many times as there are heads.

Why Does This Exist?

Multi-head attention runs H parallel attention heads, and each one keeps its own keys and values. During training that's fine — you compute one big forward pass, clear the activations, move on. During inference it's expensive in a different way: every token generated so far has to be kept around so future tokens can attend to it, and the cost is $H \times \text{seq\_len} \times d_{head}$ per layer, per generation step. That tally grows with every token you generate, and for a model with 32 heads and a long context it becomes the thing that limits how many parallel requests you can serve simultaneously.

Noam Shazeer introduced multi-query attention (MQA) in 2019 as a targeted fix. The insight is simple: heads don't need their own K and V projections to do useful work. Queries still project separately per head, keeping the expressive diversity that multiple heads buy. But keys and values use a single shared projection. The KV cache drops from $H$ copies to one copy — a factor-of-H reduction in inference memory.

Think of It Like This

One reference book, many readers

Imagine 32 researchers all reading the same reference book to answer their individual questions. Each researcher has their own question (query), and they're all consulting different sections of the same document (shared keys and values). None of them needs their own private copy of the book — only their own bookmark for what to look up. The book is shared; the search is independent.

Multi-query attention works the same way. Each attention head forms its own query from the current token, then all heads look up the same keys and values. Shared lookup, per-head question.

How It Actually Works

The single change to the projection

In standard multi-head attention, the Q, K, and V projections are all of shape $[d_{model}, H \cdot d_{head}]$ , then split into H slices. MQA keeps the query projection exactly the same but replaces the K and V projections with a single pair of shape $[d_{model}, d_{head}]$ — one head's worth, shared by all.

At inference time, the KV cache stores one $\text{seq\_len} \times d_{head}$ tensor for K and one for V (per layer), instead of H of each. If you run a 32-head model on a 4,096-token context, that's a 32× reduction in cached bytes for those projections.

What the attention computation looks like

Each head $h$ still computes its score matrix the usual way:

\text{Attention}(Q_h, K, V) = \text{softmax}\!\left(\frac{Q_h K^\top}{\sqrt{d_{head}}}\right) V

The difference is that $K$ and $V$ are now unindexed by head — they're the same tensor for all $h$ . Each head still produces a different output because each $Q_h$ is different, which is where the representational diversity comes from.

The quality trade-off

Sharing K and V forces all heads to look through the same "lens" when deciding which keys are relevant. In practice, quality drops slightly compared to full MHA, particularly on tasks that seem to benefit from heads attending over genuinely different subspaces of the context. The tradeoff was considered acceptable for generation tasks, where the latency and memory savings are most felt.

Show Me the Code

Computing how much the KV cache shrinks, then running a small verified attention pass.

1import numpy as np2
3rng = np.random.default_rng(42)4
5seq, d_model, H, d_head = 4096, 512, 32, 16  # Mistral-style shape6
7# MHA KV cache: H × seq × d_head, two tensors (K + V)8mha_cache_bytes = 2 * H * seq * d_head * 4  # float329# MQA KV cache: 1 × seq × d_head, two tensors10mqa_cache_bytes = 2 * 1 * seq * d_head * 411
12print(f"MHA KV cache (one layer): {mha_cache_bytes / 1024:.1f} KB")13print(f"MQA KV cache (one layer): {mqa_cache_bytes / 1024:.1f} KB")14print(f"Reduction factor: {mha_cache_bytes // mqa_cache_bytes}×")15# MHA KV cache (one layer): 16384.0 KB16# MQA KV cache (one layer):   512.0 KB17# Reduction factor: 32×18
19# Tiny verified attention pass (seq=4, H=2, d_head=2)20s, h, d = 4, 2, 221Q = rng.standard_normal((h, s, d))       # per-head queries22K = rng.standard_normal((s, d))          # shared keys23V = rng.standard_normal((s, d))          # shared values24
25scores = np.einsum("hqd,kd->hqk", Q, K) / d**0.526weights = np.exp(scores - scores.max(-1, keepdims=True))27weights /= weights.sum(-1, keepdims=True)28out = np.einsum("hqk,kd->hqd", weights, V)29print("output shape:", out.shape)        # (2, 4, 2)

The reduction factor matches the number of heads exactly — 32× here, and the output shape confirms H separate attention outputs despite the shared KV.

Watch Out For

Confusing MQA quality loss with a bug

When you switch a trained MHA model to MQA inference, quality drops noticeably. That's expected — the model was never trained with shared KV heads. The benefit of MQA is only fully realised when the model is trained from scratch with the shared projection, so the heads learn to work within that constraint. Don't benchmark MQA by patching a pre-trained MHA checkpoint and measuring the regression.

Assuming MQA is the current standard

MQA was a significant step but grouped-query attention has largely replaced it in practice — it keeps a small number of KV groups (typically 8) rather than collapsing to one, recovering most of the quality while still cutting the cache. If you're deciding what to implement now, GQA is the safer choice.

The Quick Version

Multi-head attention caches $H$ separate K and V tensors per layer; MQA caches one shared pair.
Queries are still projected per head, so each head produces a different output — diversity is preserved there.
KV cache drops by a factor of $H$ , which directly translates to lower serving memory and higher throughput.
Quality falls slightly because heads share the same key-value lens; the gap is bigger if the model was pre-trained with full MHA.
Grouped-query attention is the successor that keeps a handful of groups instead of collapsing to one.