T2.3 — From MHA to GQA, MQA, and MLA¶

In one line: Modern attention variants mainly change how much key/value state is stored and shared, while preserving many query heads for model capacity.


Skill	00-foundations
Module	M2 — Transformer Internals from the Inference Angle
Audience	New graduate engineer
Time	~105 min
Prereqs	T2.2
Status	built

Why this matters¶

At long context and high concurrency, KV cache can consume more memory than the weights. Reducing KV heads can shrink that cache by 4×, 8×, or more, but the architecture—not a serving flag—decides whether that sharing is valid.

The mental model¶

Suppose 32 readers ask questions about the same document. Multi-head attention gives every reader a private index and private notes. Multi-query attention keeps 32 different questions but makes all readers share one index and one set of notes. Grouped-query attention forms small reader groups that share notes within each group.

Queries represent different ways to ask what matters. Keys and values form the stored searchable history. Sharing only K and V preserves many query perspectives while reducing the history that must be stored for every token.

Many private KV paths contrast with many queries grouped onto fewer shared KV paths.

GQA keeps many query paths while several queries share each stored K/V path.

flowchart LR
    M[MHA<br/>32 Q : 32 KV] --> G[GQA<br/>32 Q : 8 KV]
    G --> Q[MQA<br/>32 Q : 1 KV]
    G --> L[MLA<br/>compressed latent cache]

Moving right reduces cached state, but each architecture requires weights trained for its sharing rule.

The mechanism¶

Let Hq be query heads and Hkv be KV heads. Standard multi-head attention has Hq=Hkv; every query head uses its corresponding K and V head. Grouped-query attention has Hkv < Hq and requires Hq to divide into KV groups. If 32 query heads share eight KV heads, four query heads use each K/V pair. Multi-query attention is the endpoint Hkv=1.

The cache formula from T2.2 contains Hkv but not Hq:

\[ \text{KV bytes}=2LBT H_{kv}d_hs. \]

That is the systems payoff. For the artifact’s example—32 layers, batch 16, 4,096 tokens, head dimension 128, bf16—MHA with 32 KV heads uses 32 GiB. GQA with eight uses 8 GiB. MQA uses 1 GiB. The reductions are exactly proportional to KV-head count because every other term remains fixed.

MHA, GQA, MQA, and latent cache stacks decrease in stored width.

Fewer KV heads shrink each token’s cache footprint in direct proportion.

Projection parameters shrink too. Queries still project from d_model to Hq × d_h, normally the full hidden width. K and V project only to Hkv × d_h. For a 4096-wide model with 32 query heads and eight KV heads, Q is 4096×4096, while K and V are each 4096×1024. The output projection remains approximately 4096×4096.

Trace one group to make the layout concrete. Query heads 0–3 may use KV head 0, query heads 4–7 use KV head 1, and so on. Each query still computes different scores because its Q vector differs. What is shared is the cached set of keys it searches and values it combines. GQA does not collapse the four query outputs into one head.

Implementations may broadcast KV heads to match query groups during the attention calculation. That logical expansion should not duplicate the stored cache. The physical cache retains only Hkv heads; otherwise the capacity advantage would disappear.

An MHA checkpoint learned separate K and V projection weights for every head. Deleting three-quarters of those heads changes the function. You can train or convert a model using a quality-preserving procedure, but setting num_key_value_heads in a JSON file does not make existing tensor shapes and learned behavior compatible.

Head sharing must match trained projection shapes; a configuration edit cannot create it.

The architecture makes a quality–capacity trade. More independent KV heads give attention more ways to represent stored content. Sharing reduces this freedom but often retains much of the quality when included during training or careful uptraining. The correct comparison uses task evaluations, not a theoretical assertion that more or fewer heads must always win.

MLA compresses a different object¶

Multi-head latent attention does not merely choose a smaller number of ordinary KV heads. It projects token state into a compressed latent representation that can be cached, then reconstructs or absorbs required projections during attention. RoPE-related components may be handled separately. The exact tensors and formulas depend on the model implementation.

For serving, the question is still familiar: what bytes are stored per token and what computation reconstructs the attention inputs? Do not estimate MLA cache by inserting a guessed “KV head count” into a GQA formula. Read the model’s latent dimensions and cache implementation.

These variants also interact with tensor parallelism. Query and KV head counts must shard across the chosen TP degree, or the engine must replicate some state. Eight KV heads divide naturally across TP widths 1, 2, 4, and 8. A wider TP layout may require replication or a specialized kernel. Mathematical validity does not guarantee an efficient engine layout.

In practice¶

Read num_attention_heads, num_key_value_heads, and head_dim from configuration. If KV heads are missing, inspect model source before assuming MHA; some families derive or rename the field. Confirm checkpoint projection shapes because configurations can be incomplete or custom code can override them.

Use the cache sweep to build intuition, then size real deployments with concurrency and length distributions. Cache savings may enable a larger batch, which improves decode weight reuse and throughput. The indirect throughput gain can be more valuable than the small reduction in projection parameters.

Record quality evidence separately from capacity evidence. The chart proves a byte ratio. A benchmark suite proves whether a trained variant retains acceptable behavior. Combining them into one undocumented “best attention” ranking conceals the actual decision.

A useful comparison records model revision, query/KV heads, cached bytes per token, concurrency, TTFT, TPOT, and task scores. Keep weights and workload fixed where possible. Comparing unrelated MHA and GQA models cannot isolate head sharing because training data, parameter count, and optimization also changed.

Failure modes¶

Symptom	Cause	Fix
Checkpoint fails to load after editing KV-head count	Projection tensor shapes no longer match	Restore configuration or use weights trained/converted for that architecture
Cache estimate does not shrink with GQA	Query heads were used in the formula instead of KV heads	Use `num_key_value_heads` for K/V storage
GQA cache ratio is correct but throughput barely changes	Another resource—weights, compute, scheduler, or communication—binds performance	Measure the full serving workload and achieved batch
MLA estimate is nonsensical	A GQA head formula was applied to latent attention	Inspect cached latent and positional tensors in model source
Quality comparison contradicts a public benchmark	Model, conversion, tasks, context, or decoding settings differ	Reproduce on identical weights and evaluation conditions

Do it¶

Run the KV-head sweep. Before execution, calculate cache for one, eight, and 32 KV heads. Then choose a real configuration and explain its sharing ratio from projection shapes.

Success means predicting every plotted capacity point exactly and explicitly labeling the chart as capacity evidence, not quality evidence.

Check¶

With 32 query heads and eight KV heads, how many query heads share each K/V head?
Which terms change in the KV-cache formula when MHA becomes GQA?
Why can you not safely convert an MHA checkpoint by editing only num_key_value_heads?