Multi-Query Attention.

Article 2.1 established the problem: the KV cache grows linearly with sequence length, and at 128K context it can outweigh the model itself. The obvious engineering question is — what's the minimum information we actually need to store?

Multi-Query Attention (MQA) was Noam Shazeer's answer in 2019. The paper is titled "Fast Transformer Decoding: One Write-Head is All You Need" (arXiv:1911.02150). The title is a deliberate riff on "Attention Is All You Need" — and the claim is as bold as the original. You can share a single key-value head across all query heads and barely lose anything in quality.

MQA is not a new idea that appeared in DeepSeek. It's a 2019 paper that was largely ignored, then adopted by PaLM (2022), Falcon, and StarCoder when serving costs became real. Understanding MQA is step 1 in understanding why GQA and MLA exist.

The MHA baseline

Standard Multi-Head Attention (MHA) from Phase 1 has H query heads, H key heads, and H value heads. Each head operates in its own d_h = d_model / H dimensional subspace.

For a model with d_model = 4096 and H = 32 heads:

• Each head dimension: d_h = 128
• Query projection: W_Q ∈ R^{4096 × 4096} (32 heads × 128 dims)
• Key projection: W_K ∈ R^{4096 × 4096}
• Value projection: W_V ∈ R^{4096 × 4096}

The KV cache stores all 32 key heads and all 32 value heads. For every token in the context, at every layer, you're storing 64 vectors of dimension 128. That's what you calculated in 2.1.

The MQA insight

Shazeer's observation: the query heads benefit from specialisation. You want 32 different questions being asked of the context — that's the diversity that makes multi-head attention powerful. But the keys and values? They're just answering those questions. You don't necessarily need a different key-value pair for each head.

MQA keeps H query heads but collapses to a single K head and a single V head, shared across all H queries:

• Query projection: W_Q ∈ R^{d_model × (H × d_h)} — same as MHA
• Key projection: W_K ∈ R^{d_model × d_h} — one head only
• Value projection: W_V ∈ R^{d_model × d_h} — one head only

The KV cache now stores 1 key head and 1 value head instead of H of each. At H=32, that's a 32× reduction in KV cache memory.

The attention computation with shared KV

In MHA, head i computes:

head_i = softmax(Q_i @ K_i^T / sqrt(d_h)) @ V_i

In MQA, all heads share the same K and V:

head_i = softmax(Q_i @ K_shared^T / sqrt(d_h)) @ V_shared

The output of each head still has dimension d_h. They still get concatenated and projected by W_O. The interface to the rest of the transformer is identical. The only change is inside the attention computation: queries remain diverse, keys and values are singular.

In PyTorch terms, the difference is surgical:

# MHA
W_Q = nn.Linear(d_model, H * d_h, bias=False)
W_K = nn.Linear(d_model, H * d_h, bias=False)  # H heads
W_V = nn.Linear(d_model, H * d_h, bias=False)  # H heads

# MQA — only K and V change
W_Q = nn.Linear(d_model, H * d_h, bias=False)
W_K = nn.Linear(d_model, d_h, bias=False)       # 1 head
W_V = nn.Linear(d_model, d_h, bias=False)       # 1 head

During the forward pass, the single K and V head gets broadcast across all H queries before the dot product:

def mqa_forward(x, W_Q, W_K, W_V, W_O, H, d_h):
    B, S, D = x.shape

    # Q: (B, S, H, d_h)
    Q = W_Q(x).reshape(B, S, H, d_h).transpose(1, 2)

    # K, V: (B, S, d_h) → broadcast to (B, 1, S, d_h) → (B, H, S, d_h)
    K = W_K(x).reshape(B, S, 1, d_h).transpose(1, 2).expand(B, H, S, d_h)
    V = W_V(x).reshape(B, S, 1, d_h).transpose(1, 2).expand(B, H, S, d_h)

    # Standard scaled dot-product attention from here
    scale = d_h ** -0.5
    attn = (Q @ K.transpose(-2, -1)) * scale
    attn = attn.softmax(dim=-1)
    out = attn @ V  # (B, H, S, d_h)

    out = out.transpose(1, 2).reshape(B, S, H * d_h)
    return W_O(out)

The KV cache stores only K_shared and V_shared — one tensor per layer, not H. At decode step t, you append one row to each. The broadcast to H query heads happens at attention computation time, not at storage time.

Memory reduction: the numbers

Applying the formula from 2.1 with H_kv = 1 instead of H_kv = H:

For a model with H=32 heads, same 70B-class architecture:

• MHA at 128K context: ~40 GB KV cache
• MQA at 128K context: ~40 GB / 32 = ~1.25 GB KV cache

That's not a typo. 32× reduction in KV cache memory. The entire 128K context fits in the cache budget that MHA would spend on a 4K context.

This is why PaLM used MQA when they scaled to 540B parameters and trillion-token context. At that scale, the KV cache isn't just a concern — it's the binding constraint on how many users you can serve simultaneously.

What you actually lose

Shazeer's paper measured quality degradation and found it acceptable for autoregressive language modelling. But "acceptable" deserves scrutiny.

The original MQA paper reported perplexity degradation of roughly 0.1-0.3 perplexity points on language modelling benchmarks — statistically significant but not catastrophic. Later studies (and the GQA paper from Google, 2023) found that the degradation is task-dependent. On tasks that require precise cross-head information aggregation — complex reasoning, structured generation — MQA's quality loss is more noticeable.

The intuition: when all H query heads attend to the same K and V, they can no longer develop fully independent attention patterns. Head 1 might want to focus on syntactic dependencies while head 8 focuses on coreference — but they're both asking slightly different questions of the same keys. With MHA, each head has its own key projections, so the question-answer pairings are truly independent. MQA forces all queries to look at the same "representation" of the context, just from different query angles.

In practice, this manifests as degraded performance on tasks requiring diverse, independent attention patterns. Factual recall, code generation, and mathematical reasoning all show measurable but small regressions in systematic benchmarks comparing MHA-equivalent and MQA-equivalent models of the same size.

Training tricks for quality recovery

Shazeer's paper also noted that MQA benefits from somewhat different training hyperparameters — specifically, the K and V projections should be initialized and tuned carefully because they carry more load (all H query heads are counting on them). Some practitioners find that increasing the K and V hidden dimension slightly (e.g., d_h × 1.5 for the shared heads) partially recovers quality at modest memory cost.

The other approach is converting existing MHA models to MQA. The GQA paper showed that you can take a fully trained MHA checkpoint and upcast it to MQA by mean-pooling the H key and value heads into one. Then fine-tune for a small number of steps to recover performance. This is cheaper than training MQA from scratch and enables you to experiment with the quality tradeoff without full retraining.

# Mean-pool H KV heads into 1 (for checkpoint conversion)
# W_K_mha: (H * d_h, d_model) → reshape to (H, d_h, d_model) → mean over H
W_K_mha = checkpoint["W_K"].reshape(H, d_h, d_model)
W_K_mqa = W_K_mha.mean(dim=0)  # (d_h, d_model)

W_V_mha = checkpoint["W_V"].reshape(H, d_h, d_model)
W_V_mqa = W_V_mha.mean(dim=0)  # (d_h, d_model)

Where MQA is used

MQA appears in several major models:

Model	Released	Architecture	Why MQA
PaLM	2022	MQA	540B model, trillion tokens — KV cache was prohibitive
Falcon-7B/40B	2023	MQA	Inference efficiency, commercial serving
StarCoder	2023	MQA	Fast inference for code completion use case
GPT-J, CodeGen	2021-22	MHA	Pre-MQA awareness; full KV heads
LLaMA-2/3	2023-24	GQA (not MQA)	G=8 groups — quality-memory compromise

Notice that the LLaMA family doesn't use MQA — it uses GQA with G=8 groups. That's the story in the next article: the ML community found that MQA was too aggressive a compression, and GQA (Grouped-Query Attention) offers a better tradeoff.

The bandwidth argument for MQA

Beyond raw memory, MQA has a bandwidth advantage that's easy to overlook. During decode, the GPU must load the KV cache from HBM on every step. With MQA, you're loading 1/H as much data. For an H=32 model, the decode step goes from loading 40 GB (at 128K context) to loading 1.25 GB. At A100's 2 TB/s bandwidth:

• MHA: 40 GB / 2 TB/s = 20 ms just to load the KV cache
• MQA: 1.25 GB / 2 TB/s = 0.625 ms for the same operation

That's a 32× bandwidth saving translating to 32× faster KV cache reads per decode step. This is why MQA wasn't just about fitting models on fewer GPUs — it was about making long-context generation usably fast.

The KV cache bottleneck is a bandwidth problem, not just a memory problem. MQA solves both simultaneously: less data to store, less data to load on every decode step.

The limitation MQA left unsolved

MQA's 32× compression comes with a quality floor. For models being deployed as general-purpose assistants or in reasoning-heavy applications, the quality regression at extreme compression ratios is unacceptable. The practitioners needed something in between.

Grouped-Query Attention (GQA), published by Google Research in 2023, provides exactly that. Instead of collapsing to 1 KV head, GQA uses G groups — each group shares a K and V head, and G can be tuned to trade off memory against quality. LLaMA-2 uses G=8, which means 8 shared KV heads instead of 32 (4× compression) and 1 (32× compression). Article 2.3 covers how this works and why it became the dominant approach.