MHA vs. MQA vs. GQA: Attention Architecture Tradeoffs

Lab: Run the companion CPU exercise at /labs/track1/03-cache-and-gqa/ (KV cache ratios, head grouping, architecture choice).

When inference is the bottleneck, the attention mechanism is not only a quality design — it is a memory-bandwidth design. During autoregressive decoding, every newly generated token requires streaming the accumulated key-value (KV) history from GPU High Bandwidth Memory (HBM). Under standard Multi-Head Attention (MHA), each query head maintains its own dedicated key and value heads, causing the KV cache to expand rapidly with context length, batch size, and layer count.

For the same layer count, head dimension, context length, batch size, and cache precision, changing from 32 KV heads to 8 reduces raw KV cache storage to one quarter.

For real architecture examples, compare Meta’s original Llama 2 7B model (which used standard MHA with 32 KV heads; Touvron et al. 2023, Table 1) and Llama 3 8B (which adopted Grouped-Query Attention with 8 KV heads; Dubey et al. 2024, §2.1; Hugging Face Meta-Llama-3-8B/config.json). Both models share 32 query heads (H_q = 32) and head dimension d_head = 128, but Llama 3 groups 4 query heads per KV head (H_kv = 8), implementing the grouped-query attention framework established by Ainslie et al. (2023), §2).

The opening memory calculation: 32 versus 8 KV heads

Hold the following Llama-inspired dimensions fixed and change only the KV head count. This calculated example compares cache storage; it does not establish equal model quality:

• Transformer layers (L): 32
• Head dimension (d_head): 128 (d_model = 4,096, H_q = 32 query heads)
• Batch size (B): 4 concurrent sequences
• Context length (S): 4,096 tokens per sequence
• Cache precision (b): 2 bytes per element (FP16 / BF16)

Using the raw KV cache sizing formula from Lesson 13:

KV bytes = 2 × L × H_kv × d_head × S × B × b

(where the factor of 2 accounts for separate Key and Value representations)

• Baseline MHA (H_kv = 32 KV heads): 2 × 32 × 32 × 128 × 4,096 × 4 × 2 = 8,589,934,592 bytes
  ≈ 8.59 GB (8.00 GiB)
• Grouped-Query Attention (H_kv = 8 KV heads, 4:1 grouping): 2 × 32 × 8 × 128 × 4,096 × 4 × 2 = 2,147,483,648 bytes
  ≈ 2.15 GB (2.00 GiB)
• Multi-Query Attention (H_kv = 1 KV head, all 32 query heads share 1 KV head; Shazeer, 2019): 2 × 32 × 1 × 128 × 4,096 × 4 × 2 = 268,435,456 bytes
  ≈ 0.27 GB (0.25 GiB)

Changing from 32 KV heads to 8 frees up 6.00 GiB (6.44 GB) of GPU VRAM on just four concurrent 4,096-token streams. In a serving cluster, extra cache headroom may support longer sequences or higher concurrency if other memory and performance constraints permit.

Why cache savings do not guarantee equivalent latency or throughput gains

A common misconception is that reducing KV cache storage to one quarter (4× reduction) automatically yields a 4× reduction in generation latency or a 4× boost in throughput. In production systems, it does not:

1. Other operations remain: Holding other dimensions fixed, fewer KV heads reduce the key/value projection sizes and cache storage. Query/output projections and feed-forward operations still contribute to execution time. Comparing Llama 2 and Llama 3 also involves other architecture and training differences, so their performance cannot be attributed to GQA alone.
2. TTFT includes more than GPU prefill: Client-observed time to first token includes queueing, scheduling, prompt processing, and transport. Prefill bottlenecks depend on workload, kernels, and hardware. Measure these components under the intended load before predicting latency.
3. Stored bytes differ from transferred bytes: A dense model with approximately 8 billion FP16 parameters has a weight-storage budget of 8 billion × 2 bytes ≈ 16 GB, excluding cache and runtime allocations. Decode batch steps can reuse weights across sequences. Actual HBM traffic also depends on kernels and cache reuse; a 4× raw KV storage reduction is not a guaranteed 4× attention-traffic or latency reduction.

Checkable understanding check

Question: An engineer deploys an 8B model (L = 32, d_head = 128, FP16) serving batch size B = 8 with context length S = 2,048 tokens.

1. Calculate the exact total KV cache footprint in gigabytes (GB, 10^9 bytes) and gibibytes (GiB, 2^30 bytes) under MHA (H_kv = 32) versus GQA (H_kv = 8).
2. Does this cache calculation establish a fourfold reduction in client-observed TTFT? Explain why or why not.

Answer:

1. Cache Footprint:
   • MHA (H_kv = 32):
     2 × 32 × 32 × 128 × 2,048 × 8 × 2 = 8,589,934,592 bytes
     ≈ 8.59 GB (8.00 GiB).
   • GQA (H_kv = 8):
     2 × 32 × 8 × 128 × 2,048 × 8 × 2 = 2,147,483,648 bytes
     ≈ 2.15 GB (2.00 GiB).
   • Reduction: Storage drops to exactly one quarter (2,147,483,648 / 8,589,934,592 = 8 / 32 = 0.25), saving 6.00 GiB (6.44 GB).
2. TTFT Latency Impact: No. The calculation establishes a cache-storage ratio only. Other model operations remain, and client TTFT includes queueing, scheduling, prefill, and transport. Their relative costs can change under load. A latency claim needs a benchmark with the same workload and measurement boundaries.

We assume the projection mechanics from lesson 6 are already in place. We do not re-derive W_in, head splitting, or contiguous(); instead we show exactly how the K/V projection shape changes, what the cache stores, and what can go wrong when grouping is implemented incorrectly.

1. What changes: query heads versus KV heads

In lesson 6, the number of attention heads h was the same for Q, K, and V. We now separate them:

• H_q: number of query heads (the h from lesson 6).
• H_kv: number of key/value heads. Each K/V head is shared across some subset of query heads.

The three architectures:

Variant	H_q (query heads)	H_kv (KV heads)	Sharing rule
MHA	4	4	Each Q head has its own K and V head (1:1).
MQA	4	1	One K head and one V head shared by all Q heads.
GQA	4	2	Q heads divided into 2 groups of 2; each group shares one K/V head.

In GQA, H_q must be divisible by H_kv for an even grouping. Here 4 / 2 = 2 query heads per KV head. A non-divisible split (e.g. H_q = 5, H_kv = 2) requires either unequal groups or a different grouping strategy; the original GQA paper uses mean-pooling of the original MHA heads to form each group head (Ainslie et al. 2023).

Head grouping: MHA, GQA, MQA

Left to right: MHA (4 KV heads, one per Q); GQA (2 KV heads, each shared by 2 Q heads); MQA (1 KV head shared by all 4 Q heads). Relative cache bars below.

Read this diagram as text

Three panels. Panel 1 (MHA): Q0, Q1, Q2, Q3 each connect to their own K/V head; 4 KV heads total; full-size cache bar. Panel 2 (GQA): Q0+Q1 share one K/V head, Q2+Q3 share another; 2 KV heads total; half-size cache bar. Panel 3 (MQA): Q0, Q1, Q2, Q3 all connect to one K/V head; 1 KV head total; quarter-size cache bar. Below each panel: relative KV-cache size annotation.

2. Trace Q, K, V and cache tensor shapes

We use the same numerical framework as lesson 6 (B = 2, T = 5, d_head = 8, d = 32) and add H_q = 4. We keep the dimensions identical so the comparison is direct.

2.1 Projection and split (same for all variants)

The Q projection produces H_q heads:

Q = X @ W_Q            # [B, T, H_q * d_head] = [2, 5, 32]
Q = Q.view(B, T, H_q, d_head).transpose(1, 2)  # [B, H_q, T, d_head] = [2, 4, 5, 8]

For MHA, K and V projections have the same shape as Q. For MQA and GQA, the K/V projections have fewer heads:

# MQA: one K head, one V head
K = X @ W_K            # [B, T, H_kv * d_head] = [2, 5, 8]
K = K.view(B, T, H_kv, d_head).transpose(1, 2)  # [B, 1, 5, 8]
V = X @ W_V            # same shape as K

# GQA: two K/V heads (4 query heads / 2 groups = 2 heads per group)
K = X @ W_K            # [B, T, H_kv * d_head] = [2, 5, 16]
K = K.view(B, T, H_kv, d_head).transpose(1, 2)  # [B, 2, 5, 8]

Note: the projection weights W_K and W_V have different output dimensions (H_kv * d_head instead of H_q * d_head). The total parameter count for K/V projections drops linearly with H_kv / H_q.

2.2 Attention computation per variant

For MHA, attention runs independently for each of the 4 heads (lesson 6). For MQA and GQA, the attention scores are computed with the same Q but with fewer K heads:

# MHA: [2, 4, 5, 5] scores per head
scores = Q @ K.transpose(-2, -1) / sqrt(d_head)  # [B, H_q, T, T]

# MQA: Q [2, 4, 5, 8] attends to K [2, 1, 5, 8]
# Broadcast K across the H_q dimension before matmul, or repeat K along head dim
K_expanded = K.repeat(1, H_q // H_kv, 1, 1)  # For MQA: repeat 4 times → [2, 4, 5, 8]
scores = Q @ K_expanded.transpose(-2, -1) / sqrt(d_head)  # [2, 4, 5, 5]

In practice, frameworks handle this by either repeating K/V before the matmul or by reshaping the score computation. The architectural point is that the same K/V tensor is read by multiple query heads, not that new independent K/V projections are computed.

2.3 Cache shapes after prefill (S = 5 tokens cached)

The KV cache stores the output of the K and V projections, not the projected Q. From the KV-cache lesson (lesson 13):

# MHA cache per layer
K_cache: [B, H_kv, S, d_head] = [2, 4, 5, 8]
V_cache: [2, 4, 5, 8]

# MQA cache per layer
K_cache: [B, H_kv, S, d_head] = [2, 1, 5, 8]
V_cache: [2, 1, 5, 8]

# GQA cache per layer
K_cache: [B, H_kv, S, d_head] = [2, 2, 5, 8]
V_cache: [2, 2, 5, 8]

The query heads do not change the cache; only H_kv does.

3. Per-layer KV-cache byte derivation

From lesson 13, the raw cache formula is:

KV bytes = 2 × L × H_kv × d_head × S × B × b

Where:

• 2: one key and one value per token.
• L: number of transformer layers.
• H_kv: number of KV heads per layer (varies by variant).
• d_head: head dimension.
• S: total cached tokens (prompt + generated).
• B: batch size.
• b: bytes per element (2 for FP16, 4 for FP32).

For a single layer, L = 1:

per_layer_KV_bytes = 2 × B × S × H_kv × d_head × b

This is the formula we use for the comparison below.

3.1 Worked calculation: all three variants, same dimensions

Use B = 2, S = 5, d_head = 8, b = 2 (FP16), L = 1 (per-layer comparison).

Variant	H_kv	Per-layer K+V bytes (formula)	Per-layer total
MHA	4	2 × 2 × 5 × 4 × 8 × 2	1,280 bytes
GQA	2	2 × 2 × 5 × 2 × 8 × 2	640 bytes
MQA	1	2 × 2 × 5 × 1 × 8 × 2	320 bytes

Ratio relative to MHA:

• GQA: 640 / 1,280 = 0.50 (50% of MHA cache).
• MQA: 320 / 1,280 = 0.25 (25% of MHA cache).

These are cache-size ratios only. They do not imply equivalent reductions in total model memory (weights dominate), nor do they guarantee identical latency reductions (the query projection W_Q and attention computation over H_q heads remain unchanged). The memory-bandwidth reduction comes specifically from reading fewer K/V tensors during decode.

Cache savings are not latency guarantees

Fewer KV heads reduce the bytes read from memory per decode step. However, the query projection (W_Q) and attention computation (Q @ K^T) scale with H_q, not H_kv. A model with H_q = 32 and H_kv = 8 (GQA) still performs 32 query-head computations; only the K/V load is 4× smaller. Actual latency depends on kernel fusion, batch size, sequence length, GPU memory bandwidth, and whether the bottleneck is compute or memory (Shazeer 2019; Ainslie et al. 2023).

3.2 Scaling to full model and full context

Connecting back to our opening 32-layer benchmark (L = 32, d_head = 128, S = 4,096, B = 4, b = 2 in FP16):

per_layer_KV_bytes = 2 × B × S × H_kv × d_head × b

• MHA (H_kv = 32): Each layer stores 2 × 4 × 4,096 × 32 × 128 × 2 = 268,435,456 bytes (256 MiB ≈ 0.268 GB).
  Across all 32 layers: 32 × 256 MiB = 8,192 MiB = 8.00 GiB (8.59 GB).
• GQA (H_kv = 8): Each layer stores 2 × 4 × 4,096 × 8 × 128 × 2 = 67,108,864 bytes (64 MiB ≈ 0.067 GB).
  Across all 32 layers: 32 × 64 MiB = 2,048 MiB = 2.00 GiB (2.15 GB).
• MQA (H_kv = 1): Each layer stores 2 × 4 × 4,096 × 1 × 128 × 2 = 8,388,608 bytes (8 MiB ≈ 0.008 GB).
  Across all 32 layers: 32 × 8 MiB = 256 MiB = 0.25 GiB (0.27 GB).

The linear reduction factor holds identically at per-layer scale and full-model scale because every layer independently stores H_kv / H_q as many key and value vectors.

4. Quality tradeoff and why evaluation is required

Reducing H_kv means fewer independent key/value representations. Each KV head can capture different patterns (e.g., syntax, entity tracking, positional relations). Sharing heads forces multiple query heads to attend to the same representation, which may reduce capacity.

From the original papers:

• MQA (Shazeer 2019): “Multi-query attention … can indeed be much faster to decode, and incur only minor quality degradation from the baseline.” The paper reports results on WMT 2014 English-German translation; degradation is small but measurable.
• GQA (Ainslie et al. 2023): “Uptrained GQA achieves quality close to multi-head attention with comparable speed to MQA.” The authors convert MHA checkpoints via mean-pooling of heads within each group and uptrain with ~5% of original pretraining compute.

These claims are model-specific and task-specific. A GQA conversion for a code-generation model may show different behavior than for a translation model. The architecture decision requires:

1. A baseline MHA evaluation on your target metric.
2. A converted GQA or MQA evaluation with the same evaluation pipeline.
3. A comparison of memory-bandwidth reduction (measured, not estimated) and the observed quality delta.
4. A decision whether the latency/memory improvement justifies the quality change for the deployment context.

There is no unconditional ranking. If the workload is memory-bandwidth-bound and quality degradation is within the acceptable threshold, GQA or MQA is defensible. If quality is the overriding constraint and the workload is compute-bound, MHA remains the safer choice.

4.1 Engineering decision example

Consider a production serving scenario:

• Constraint A: decode latency must stay below 50 ms/token at batch size 8, sequence length 2K.
• Constraint B: perplexity degradation must stay within 2% of the MHA baseline on the evaluation set.
• Constraint C: no retraining budget beyond checkpoint conversion (uptraining only).

Under these constraints:

• MQA achieves the largest bandwidth reduction (32× fewer K/V heads: all 32 query heads share a single K/V head) but carries the highest quality risk; the original MQA paper reports minor degradation but the magnitude depends on model size and data.
• GQA with H_kv = 8 (one quarter of H_q = 32) achieves a 4× reduction with quality close to MHA (Ainslie et al. 2023, §4.3); this is the safer middle ground when uptraining is available.
• If uptraining is unavailable and only inference-time conversion is possible, MHA remains the only defensible option because un-uptrained GQA may suffer from incorrect head mapping.

This is an architecture decision, not a universal recommendation. The choice depends on the measured quality delta on the target task, not on the ratio alone.

5. Head grouping and implementation pitfalls

The architectural difference is small in theory but error-prone in practice. Common mistakes:

5.1 Incorrect head mapping in GQA

When H_q = 4 and H_kv = 2, the query heads 0, 1 share KV head 0, and 2, 3 share KV head 1. If the code maps query head 0 to KV head 1, the attention attends to the wrong representation. The grouping must be consistent across the projection, the attention computation, and the output projection.

A simple correctness check: after grouping, verify that the number of K/V tensors equals H_kv, and that each query head’s attention weights have shape [B, T, T] (not [B, T, T, H_kv]). Any extra dimension in the K or V tensor after projection indicates a grouping error.

5.2 Incorrect repetition or broadcast

In MQA, the single K/V head must be made available to all H_q query heads. A common error is to repeat the projection weights W_K H_q times (treating MQA as MHA with shared weights) instead of projecting once to H_kv * d_head and then broadcasting or repeating the resulting tensor. Repeating the weights changes the parameter count; broadcasting the tensor does not. Verify:

# Correct MQA projection
W_K_shape = [d, H_kv * d_head]  # [32, 8], not [32, 32]

5.3 Cache write misalignment

From lesson 13: writing at write_position + 1 instead of write_position shifts the cache by one column, causing silent errors. This is independent of grouping but is more likely to be missed when the cache dimension (H_kv) has changed and the shape assertions are not updated.

Always assert after the projection:

assert K.shape == (B, H_kv, T, d_head)
assert V.shape == (B, H_kv, T, d_head)

And after the attention:

assert scores.shape == (B, H_q, T, T)

Any mismatch in the H_kv or H_q dimension indicates a grouping, projection, or repetition error before quality evaluation begins.

6. Sources and evidence limits

• Shazeer (2019) — Original MQA proposal (arXiv:1911.02150). Shows memory-bandwidth reduction and reports minor quality degradation on WMT 2014 EN-DE. Does not claim universal applicability.
• Ainslie et al. (2023) — GQA proposal (arXiv:2305.13245, EMNLP 2023). Proposes uptraining from MHA checkpoints, introduces grouped-query with G groups (G = H_kv), and evaluates on T5-XXL and LLaMA-style architectures. Reports quality close to MHA with speed close to MQA, with G = 8 selected as a favorable middle ground.
• LLaMA-2 (Touvron et al., 2023) — Table 1 of the paper lists the 7B and 13B models with standard multi-head attention (H_q = H_kv = 32) and the 34B and 70B models with GQA and H_kv = 8 (H_q = 40 for 34B, H_q = 64 for 70B). This demonstrates real-world adoption but does not publish per-variant quality comparisons for all model sizes.
• Llama-3 (Dubey et al., 2024) — Section 2.1 (arXiv:2407.21783) and official model configurations (e.g. Meta-Llama-3-8B/config.json). Demonstrates the transition of the 8B model class from MHA (H_kv = 32) in Llama 2 to GQA (H_kv = 8) in Llama 3, standardizing the 4:1 grouping ratio (H_q = 32, H_kv = 8) across the 8B scale.

We do not invent measurements or claim that any paper’s result applies to every model. The numerical ratios in this lesson (0.50 and 0.25) are derived from the stated dimensions; the quality claims are summarized from the original papers with explicit scope (translation for MQA, T5/LLaMA-style uptraining for GQA).

7. Exit exercise

Given B = 1, S = 1,024, d_head = 64, b = 2 (FP16), compute the per-layer KV-cache size in bytes for:

1. MHA with H_q = H_kv = 32
2. GQA with H_q = 32, H_kv = 8
3. MQA with H_q = 32, H_kv = 1

Then state: your evaluation shows a 3% perplexity increase for GQA (4× cache reduction: 32 KV heads down to 8) and a 6% increase for MQA (32× reduction: 32 KV heads down to 1) relative to MHA; your serving budget requires at least a 2.5× memory-bandwidth reduction while the task tolerates at most a 5% quality loss. Which variant (if any) satisfies both constraints? Explain why the answer depends on the measured quality delta, not only the ratio.

Exercise: Cache ratios and architecture choice

1. 01Compute per-layer KV-cache bytes for MHA, GQA (H_kv = 8), and MQA with B = 1, S = 1,024, d_head = 64, b = 2 (FP16). Show the formula and the numerical result for each.

2. 02State the ratio of GQA (H_kv = 8) and MQA (H_kv = 1) cache sizes relative to MHA with H_kv = 32, and express each as a reduction factor.

3. 03Given a 3% perplexity increase for GQA (4× cache reduction: 32 KV heads to 8) and a 6% increase for MQA (32× reduction: 32 KV heads to 1) relative to MHA, a required memory-bandwidth reduction of at least 2.5×, and a tolerated quality loss of at most 5%: which variant, if any, satisfies both constraints? Justify from the measured deltas.

8. Related lessons

• Lesson 6 (multi-head-attention): projection, splitting, merging, and output projection mechanics (not repeated here).
• Lesson 13 (kv-cache-explained): per-layer cache growth, prefill/decode shape trace, and raw byte formula (referenced for the derivation in section 3).
• Lesson 20 (architecture-comparison): a later architecture decision framework that builds on the tradeoffs discussed here.

Related lessons

• LLM Internals & Architecture · Attention mechanisms from scratchMulti-Head Attention: Splitting, Merging, and Projecting14 min
• LLM Internals & Architecture · Generation, KV cache, and contextKV Cache Implementation and Memory Accounting14 min

9. Sources

Source	Relevance
Vaswani et al. (2017), “Attention Is All You Need” (arXiv:1706.03762), §3.2.2	Original multi-head attention definition; MultiHead(Q,K,V) = Concat(head_1,...,head_h)W^O.
Shazeer (2019), “Fast Transformer Decoding: One Write-Head is All You Need” (arXiv:1911.02150)	Original MQA proposal; memory-bandwidth analysis; minor quality degradation reported on WMT 2014 EN-DE.
Ainslie et al. (2023), “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints” (arXiv:2305.13245, EMNLP 2023)	GQA definition (G groups of query heads sharing one K/V); uptraining recipe; quality close to MHA, speed close to MQA; G = 8 selected as middle ground.
Touvron et al. (2023), “LLaMA 2: Open Foundation and Fine-Tuned Chat Models” (arXiv:2307.09288)	Real-world adoption: standard MHA (H_kv = H_q = 32) in 7B and 13B; GQA (H_kv = 8) in 34B and 70B.
Dubey et al. (2024), “The Llama 3 Herd of Models” (arXiv:2407.21783), §2.1; Hugging Face Meta-Llama-3-8B/config.json	Real-world GQA adoption across 8B, 70B, and 405B; transitions 8B from Llama 2’s MHA (H_kv = 32) to GQA (H_kv = 8) with d_head = 128.
PyTorch nn.Linear docs (torch.nn.Linear)	Projection weight shapes; confirms W_K output dimension is user-defined, not fixed to 3d.