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 .. h

So 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.

Advertisement

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.

Advertisement

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) × b

The leading 2 counts keys and values separately. Taking the ratio, everything cancels except the head count:

MHA / MQA = (h × d_k) / d_k = h

So 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 sequence

Four 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.