Multi-Query & Grouped-Query Attention

The previous chapter made the KV cache concrete: every active token consumes GPU memory during decode. This chapter asks how the model architecture can shrink that cache before serving tricks like paging, scheduling, or quantization are applied.

Multi-query and grouped-query attention reduce KV cache memory by sharing key and value heads across attention heads. This chapter explains why that memory saving matters for long prompts, high concurrency, and production serving cost.

Imagine a fulfillment center where every picker keeps their own personal copy of the route sheet. With five pickers, the paperwork pile is small. With sixty-four pickers and thousands of orders, the storage room overflows. The fix is simple: have small teams share one route sheet instead of giving every picker a personal copy.

That's the idea behind Grouped-Query Attention (GQA) and Multi-Query Attention (MQA). In standard decoder-side multi-head attention^[1], each query head has its own Key and Value projections, so incremental decoding keeps separate cached K/V state per head. As conversations get longer and you serve more users simultaneously, this cache can grow to hundreds of gigabytes. GQA and MQA reduce that dynamic memory cost by sharing cached data across heads. If KV memory is the admission bottleneck, that saving can increase concurrency; it is not an automatic throughput or quality guarantee.

Why KV cache memory becomes the bottleneck

As we saw in the previous article on inference mechanics, every time a model generates a token, it saves the Key and Value vectors from that step so it doesn't have to recompute them later. This KV cache is essential for performance, but it grows linearly with sequence length, and in standard Multi-Head Attention (MHA), every single head keeps its own separate copy. For big models, this adds up fast.

A tiny example to build intuition

Before we scale up to billions of parameters, let's walk through a miniature case you can count on your fingers.

Suppose a single layer has 2 attention heads and each head has a dimension of 4. For one token, we must store a Key vector and a Value vector for every head. That's 2 heads times 2 vectors (K and V) times 4 numbers each:

If you serve 8 requests in parallel and each request reaches 2,048 tokens, the cache explodes to 16 x 2,048 x 8 = 262,144 numbers for that one layer alone. Scale that to 80 layers and the pile becomes enormous. The problem isn't the math itself; it's the memory needed to hold all those numbers while the GPU streams them through High Bandwidth Memory (HBM) on every decoding step.

In standard MHA, each of the $h$ attention heads maintains its own Key ($K$) and Value ($V$) projections. The attention computation for a single layer is:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

Where $Q \in \mathbb{R}^{h \times n \times d_k}$, $K \in \mathbb{R}^{h \times n \times d_k}$, and $V \in \mathbb{R}^{h \times n \times d_k}$ for a sequence of length $n$. Each head learns a different attention pattern by attending to different aspects of the input. The KV cache stores $K$ and $V$ for all previously computed tokens so the model doesn't have to recompute them on each decoding step.

KV cache memory with standard MHA

For standard Multi-Head Attention (MHA) with $h$ heads, each with dimension $d_k$, here is how to read the cache size formula before you apply it:

• 2 counts the Key and Value tensors separately.
• $h$ is the number of attention heads.
• $d_k$ is the dimension inside each head.
• bytes per element depends on dtype: 2 bytes for FP16, 1 byte for INT8, and so on.

Multiply those four numbers and you get the cache for a single token in a single layer. For the full cache across all layers, sequence length, and batch:

$\text{KV cache per token per layer} = 2 \times h \times d_k \times \text{bytes per element}$

$\text{Total KV cache} = 2 \times L \times S \times B \times h \times d_k \times \text{bytes per element}$

Where $L$ is layers, $S$ is sequence length, and $B$ is batch size.

Common trap: $2 \times h \times d_k \times \text{bytes}$ is only the cache for one token in one layer. For a full serving estimate, you still need to multiply by layers, sequence length, and batch size.

Concrete example (72B-style decoder dimensions)

Qwen2.5-72B^[2] is a useful reference point because it uses 80 layers, 64 query heads, and a head dimension of 128. The published model uses GQA-8, not full MHA, but if a model with those same dimensions used FP16 MHA, the KV cache would be:

$\text{KV cache} = 2 \times 80 \times 4096 \times 32 \times 64 \times 128 \times 2 \text{ bytes} \approx \mathbf{344 \text{ GB}}$

(Equivalent formulation using $d_{\text{model}} = h \times d_k = 8192$: $2 \times 80 \times 4096 \times 32 \times 8192 \times 2$ bytes.)

$2$ counts K and V tensors, $80$ is layers, $4096$ is tokens per request, $32$ is batch size, $64$ is heads, $128$ is head dimension ($d_k$), and the final $2$ is FP16 bytes per element.

That's more than twice the model weights themselves (~144 GB in FP16). For long-context or high-concurrency decoding, moving that cache through HBM can dominate the incremental attention path.^[3]

The next figure uses the same 72B-style dimensions, but shifts the shape of the workload from 32 parallel 4K requests to one 128K request. Both cases contain 131,072 cached token positions, so full MHA lands in the same 344 GB range.

In an e-commerce setting, this is the difference between a full-MHA cache that doesn't fit on a single 80 GB GPU and a GQA-8 cache around 43 GB before weights, runtime buffers, and fragmentation. Whether either configuration is viable still depends on weights, quantization, hardware layout, scheduler policy, and measured traffic.

---

Multi-Query Attention (MQA)

The extreme compression

Back to the warehouse analogy: MQA takes the extreme approach. Instead of every picking lane keeping its own copy of the route notes, all lanes share one copy. This saves enormous storage, but some lanes might lose details that would have helped their specific work.

MQA^[4] uses one shared K head and one shared V head across all query heads:

• $h$ separate query projections (same as standard MHA)
• 1 shared key projection for all $h$ query heads
• 1 shared value projection for all $h$ query heads

In the head-sharing figure above, MQA is the rightmost design: every query head routes through the same cached K/V pair.

How much memory MQA actually saves

Return to our tiny example: 2 heads, dimension 4, one token. Under MHA we stored 16 numbers. Under MQA we store only 1 Key and 1 Value head, so the count drops to $2 \times 1 \times 4 = 8$ numbers. That's a 2x reduction for 2 heads. At 64 heads the reduction becomes 64x.

The memory footprint of MQA is drastically smaller because the number of attention heads ($h$) in the standard MHA formula is replaced by $1$:

$\text{MQA KV cache per token per layer} = 2 \times 1 \times d_k \times \text{bytes per element}$

Where $1$ means a single shared KV head (instead of $h$ heads), so cache per token per layer scales with $d_k$ rather than with $h \cdot d_k$.

Instead of storing separate K and V for each of the 64 heads, MQA stores just one shared K and one shared V, cutting cache by 64x. All query heads look up the same key-value pair, like every picking lane using the same route note.

This scaling is proportional to the number of query heads: with $h$ query heads and a single KV head, you get an $h$x reduction. The 8-head illustration above shows this with 8x reduction; the same principle applies at scale with 64 heads for a 64x reduction.

For the same 72B-style dimensions with batch=32 and seq=4096, MHA uses 344 GB of KV cache. MQA cuts that by 64x to ~5.4 GB.

Performance vs. quality tradeoffs in MQA

MQA can reduce quality because all heads share the same key-value representation. Query heads can still ask different questions, but they no longer get separate K/V subspaces.

Common mistake: Treating MQA quality loss as either zero or catastrophic. Real impact depends on model size and task. Shazeer's original paper motivates MQA as a serving optimization, and later GQA work shows why many larger models want more than one KV head.^[4][5]

MQA is like projecting a single telescope's view onto a screen that all astronomers in an observatory can see simultaneously. They can each ask different questions about what they see (separate queries), but they're all looking at the same image (shared K/V). Most of the time, the single view is good enough, but occasionally, one astronomer misses a detail they would have caught with their own specialized lens.

A minimal MQA implementation

Here is a basic PyTorch implementation of Multi-Query Attention. It projects the input into multiple query heads, but uses only a single shared key and value projection across all query heads. Notice how the key and value tensors are broadcast across the query heads during the attention score calculation to produce the final output.

---

Grouped-Query Attention (GQA)

The balanced compromise

MQA saves the most memory, but forcing all query heads to share one K/V pair can reduce quality on some workloads. The compromise is to split the lanes into groups, with each group sharing one route note. That's GQA: less aggressive than MQA (one note for everyone), but more memory-efficient than MHA (one note per query head).

GQA^[5] is the compromise between MHA and MQA. Instead of 1 KV head (MQA) or $h$ KV heads (MHA), use $g$ KV groups where $1 < g < h$:

$\frac{h}{g} = \text{queries per group}$

With $g$ groups and $h$ total query heads, each group of $h/g$ queries shares one K and one V. For example, Mistral 7B's 32 query heads with 8 KV groups means 4 queries share each KV pair. That's a 4:1 ratio: 4x less KV cache than MHA, but with significantly more representational diversity than MQA's single KV pair.

Think of GQA like a logistics hub: instead of every driver carrying their own copy of the delivery manifest (MHA) or one manifest for the entire fleet (MQA), you organize drivers into pods that share a manifest. Each pod can still take different routes (different queries), but it shares cached reference data.

GQA in practice

To understand how this looks in production, let's look at the configurations used by several popular models:

Beyond memory savings, GQA can make tensor-parallel serving cleaner. If num_kv_heads divides the tensor-parallel degree, each shard can own an even slice of KV heads. If it doesn't divide cleanly, runtimes often replicate some KV heads instead, which wastes memory.

Memory comparison (72B model, seq=4096)

These are per-request KV cache sizes at the given sequence length:

In Ainslie et al., intermediate group counts recovered much of MHA quality while retaining MQA-like inference benefits after uptraining.^[5] The correct group count for another model remains an architecture and evaluation choice, not a guarantee inherited from that experiment.

Real-world adoption across model families

Different model families have made distinct architectural choices based on their serving requirements and quality targets:

Key observations

• MQA is best understood as a serving optimization: shared KV heads cut cache size and reduce memory traffic during incremental decoding.^[4][3]
• Many modern decoder-only LLMs use grouped KV heads, but the choice is model-specific rather than family-wide. Published examples in this article include Llama 2 34B/70B, Mistral 7B, Qwen2.5-72B, and Gemma 2 9B.^[6][7][2][8]
• There's no universal best group count. Model size, quality target, and serving stack decide whether 2:1, 4:1, or 8:1 is right.^[5]

---

Converting MHA to GQA via uptraining

A key practical insight for teams adapting older models: you don't need to train a GQA architecture from scratch. You can convert an existing MHA model to GQA through a process called "uptraining."

Key insight: Converting MHA to GQA is like a company merger. Several departments each had their own filing system (KV heads). Instead of rebuilding from scratch, you merge departments into groups, average their files, and then give the new organization time to adapt through continued training. Quality must be evaluated after the conversion.

The conversion has two main steps:

1. Mean-pool KV heads: Group the existing Key and Value heads into the desired number of partitions (e.g. merging 64 heads into 8 groups of 8). Take the mean of the weights for the $K$ and $V$ projection matrices within each group. This initializes the new GQA heads to the average behavior of the original MHA heads, preserving the model's existing attention patterns as a strong starting point.
2. Fine-tune (uptrain): Train the converted model for a small fraction of the original pretraining budget using the standard next-token prediction objective. This allows the model to adapt its internal representations to the newly merged KV projections.

The original GQA paper^[5] showed that this recipe recovers most of the lost quality with about 5% of the original pretraining compute. For a fresh pretraining run, teams usually bake the desired KV-head pattern into the architecture from day one. Uptraining matters when you inherit an older MHA checkpoint and want serving gains without restarting pretraining.

---

Beyond GQA: Multi-Head Latent Attention (MLA)

GQA reduces the number of distinct KV heads stored in the cache. Multi-Head Latent Attention (MLA) takes a fundamentally different route: it reduces the dimensionality of what must be stored per position by learning a low-rank latent representation.

How MLA works (high-level)

Instead of caching full per-head Key and Value content vectors for every token, the model does the following:

1. Input hidden states are projected into a compact latent vector of dimension $d_c$.

2. Only this compact latent state (plus a small amount of decoupled positional information) is written into the KV cache.

3. The architecture defines learned up-projections from that latent state. In an optimized inference implementation, those projection matrices can be absorbed into the query and output paths so decode does not materialize full per-head cached K/V again.^[9]

DeepSeek-V2 also separates a positional RoPE component from compressed content so the cache can retain the required position-dependent term without preventing content compression.^[9]

DeepSeek-V2 reports a 93.3% reduction in KV cache memory footprint compared to a comparable 67B baseline model, alongside its reported benchmark results.^[9] The savings are not a simple head-count ratio; they depend on the chosen latent width $d_c$, positional component, and execution path.

Trade-offs and adoption

Why the runtime distinction matters. GQA slots into grouped-attention implementations and familiar tensor-parallel strategies. MLA changes the cached representation and needs a compatible execution path. A smaller cache on paper is only useful when the selected runtime implements that path efficiently.

This architectural fork appears in production interviews: "When would you choose GQA over a more aggressive compression scheme like MLA?" The defensible answer is that GQA has a simpler K/V-head contract, while MLA can compress harder if the selected model and runtime both support its latent path.

Using DeepSeek-V2's documented 512-dimensional compressed K/V latent and 64-dimensional decoupled key component, the raw cached-value comparison looks like this.^[9]

---

Real-world serving impact

Concurrency impact (72B-style dimensions, 4K context)

Reducing KV cache doesn't guarantee an equal throughput gain, but it does raise the memory-limited concurrency ceiling. Using the same 80-layer, 64-query-head, 128-dim-head example:

If KV cache is what limits batch size, those ratios are first-order good estimates. Real systems land below that theoretical ceiling because weights, scheduler overhead, kernel efficiency, and interconnect traffic still matter.^[3]

Long-context impact

KV savings become larger in absolute bytes with long-context models. For the same 72B-style dimensions, per-request cache grows linearly with sequence length:

At long context, head-sharing or another cache compression strategy becomes an explicit capacity decision. At 128K, full MHA uses roughly 344 GB of KV state per request in this example, before weights, allocator fragmentation, or extra concurrency. GQA-8 drops that cache to about 43 GB, which is still expensive and still needs an end-to-end memory budget.

Common mistake: Treating long-context support as an architectural context-window claim rather than a serving budget. A long support history, codebase prompt, or retrieved document bundle can consume the same KV allocation. Calculate active-token memory before promising concurrency.

---

How serving engines handle GQA head expansion

When building or working with inference engines, you must handle the KV head expansion efficiently. Here is a conceptual implementation of how an engine matches the smaller number of KV heads to the larger number of query heads before computing attention.

Modern kernels (FlashAttention^[10], FlashInfer^[11]) handle GQA without materializing the expanded K/V tensors. This avoids the memory overhead of repeat_interleave and computes grouped attention directly in the fused kernel.

Production tip: When evaluating models for serving, total parameter count isn't enough. num_kv_heads, context length, and KV-cache dtype often dominate concurrency. GQA-8 raises the memory-limited concurrency ceiling by roughly 8x versus same-dimension MHA, but realized cost per query still depends on batching, quantization, and kernels.

KV cache update correctness

In serving code, the hardest GQA bug is often not the attention formula. It's updating the cache correctly one token at a time.

For one decode step, the runtime should:

1. Project only the new token into Q, K, and V.
2. Apply RoPE or another positional transform using the absolute position of that token.
3. Append the new K/V slice at the next cache position.
4. Run attention with Q from the new token against all cached K/V positions.
5. Track cache length per request, because different requests finish at different times.

The cache shape should use KV heads, not query heads:

For an MHA model, n_kv_heads == n_heads. For MQA, n_kv_heads == 1. For GQA, 1 < n_kv_heads < n_heads.

Common cache-update bugs look like this:

That's why cache tests should compare a cached decode path against a full-prefix recompute path on the same tiny prompt. The logits should match closely for the next token. If they don't, the bug is usually position IDs, mask shape, or cache append order.

---

What you should be able to defend

By this point, you should be able to reason from the architecture diagram to serving cost, not just name the attention variant.

Follow-up questions

A 70B-style chat model must serve 128K requests. What memory story do you defend?

For 80 layers, 64 query heads, head dimension 128, FP16 cache, and full MHA, the per-token-per-layer cache is 2 * 64 * 128 * 2 = 32 KB. Across all 80 layers, that is about 2.5 MiB per token. At 128K context, one request needs about 344 GB of KV cache under full MHA. GQA-8 cuts that to about 43 GB per request, which is much better but still too large for comfortable single-node serving. Strong answers mention head sharing plus paging, scheduling, retrieval, or KV-cache quantization.

A team wants GQA but cannot afford full retraining. What rollout plan makes sense?

Use uptraining. Group old K/V heads, mean-pool each group into one new K head and one new V head, then continue language-model training for a small fraction of original pretraining budget so the model adapts to reduced K/V capacity. After that, validate both task quality and serving behavior: perplexity or downstream evals, KV-cache size, tensor-parallel sharding fit, and decode throughput under realistic batch and context settings. Ainslie et al. used this path to recover much of the quality lost by reducing KV heads.^[5]

num_key_value_heads = 8 and tensor parallel degree is 6. Why is that awkward?

GQA is easiest when KV heads divide cleanly across tensor-parallel shards. With 8 KV heads and TP degree 6, some runtimes must replicate KV data or use uneven ownership. That wastes memory and cuts into the concurrency gain you expected from GQA.

Your runtime offers both GQA and MLA paths today. How do you choose between them?

Choose GQA when your selected runtime supports grouped heads and its sharding layout fits the model. Choose MLA when your model is MLA-native, your stack ships its latent-cache path, and cache pressure is dominant enough to justify that execution contract.

Cached decode diverges from full recompute after 200 tokens. What do you test first?

Check that cache shape uses n_kv_heads, not query heads, and that absolute positions, cache cursors, and append order all stay correct. The quickest proof is a tiny-prompt parity test: cached decode logits should match full-prefix recompute closely at each step.

---

Common pitfalls

"Model weights are the same thing as KV cache"

Symptom: You calculate that a quantized model fits in GPU memory, then the server still runs out of memory under long context or high concurrency.

Cause: Weights are static model parameters. The KV cache is dynamic per-request state. GQA mainly shrinks dynamic K/V activations, not the feed-forward layers or the full weight tensor.

Fix: Budget weights and KV cache separately. Then add activation buffers, allocator slack, scheduler state, and fragmentation.

"GQA speeds up training"

Symptom: You benchmark training throughput and see almost no change after switching from MHA to GQA.

Cause: Training often has a different bottleneck from incremental decoding. GQA reduces decode cache state, while training processes whole sequences and may remain dominated by large attention and feed-forward computations.

Fix: Measure decode throughput (tokens per second during autoregressive generation), not training throughput. GQA's wins show up when you're repeatedly loading the KV cache for each new token, not when you're computing the full attention matrix once.

"MQA always destroys quality"

Symptom: You reject MQA outright even for small, latency-sensitive models where the quality difference isn't visible on your task.

Cause: MQA imposes the strongest sharing constraint, but impact depends on model size, task, and training recipe. Treating it as always catastrophic is as wrong as treating it as free.

Fix: Compare task evals and serving metrics for your workload. Use GQA when you need a safer quality/serving compromise.

"I saved 8x on KV cache, so my serving cost dropped 8x"

Symptom: You calculate a huge KV cache reduction, but real-world latency or cost only improves by a fraction of that.

Cause: KV cache is one piece of the puzzle. Model weights, FFN compute, attention arithmetic, scheduler overhead, and interconnect traffic still matter. If your batch was previously limited by compute rather than memory, shrinking the cache won't move the needle as much.

Fix: Profile end-to-end latency with a realistic batch size. Use a tool like NVIDIA Nsight Systems or vLLM's built-in metrics to see whether you're memory-bound or compute-bound before betting on GQA alone.

"repeat_interleave exploded my memory during GQA training"

Symptom: Out-of-memory errors when you naively expand KV heads with repeat_interleave inside a training loop.

Cause: repeat_interleave materializes a larger tensor in memory. For 32 query heads and 8 KV groups, that temporarily creates a 4x larger K/V tensor before the matmul.

Fix: Use FlashAttention or FlashInfer, which handle GQA natively without materializing the expanded tensor. If you must write a custom kernel, use an implicit broadcast or fused kernel rather than explicit expansion.

"I can regroup heads any way I want during GQA conversion"

Symptom: The converted checkpoint loads, but quality drops sharply or tensor-parallel shards disagree after serving rollout.

Cause: Head grouping is not arbitrary. Query and KV heads often follow a fixed ordering in the weight matrices, and tensor-parallel sharding can interleave or chunk that order in specific ways. Regrouping without respecting the original layout silently changes which heads share a K/V projection.

Fix: Preserve the model's published head ordering during regrouping, then run a tiny parity test and shard-level smoke test before broader evals. Don't assume head 0-7 is always the intended first group.

"The KV-head count doesn't affect sharding"

Symptom: Tensor-parallel serving uses more memory than expected, or some shards replicate K/V state.

Cause: num_key_value_heads interacts with the tensor-parallel degree. Awkward divisibility can force replication or uneven ownership in some runtimes.

Fix: Check model config before choosing TP degree. Prefer serving layouts where KV heads divide cleanly across shards.

---

A quick check

Try this without looking back at the tables:

A model has 40 layers, 32 query heads, 8 KV groups, head dimension 128, and FP16 KV-cache elements. You serve a batch of 8 requests, each at 2,048 tokens. How large is the KV cache under GQA? Show your work.

Solution sketch

1. Per token per layer: 2 (K and V) x 8 (KV groups) x 128 (head dim) x 2 (FP16 bytes) = 4,096 bytes.
2. Per request: 2,048 tokens x 40 layers x 4,096 bytes = 335,544,320 bytes ≈ 320 MiB.
3. Batch of 8: 8 x 320 MiB ≈ 2.5 GiB.

Compare to MHA: 32 KV heads instead of 8 gives a 4x larger cache, so the same batch would need roughly 10 GiB just for the KV cache.

---

What to remember and where to go next

1. KV Cache Bottleneck: For long contexts or high concurrency, KV cache can consume more memory than model weights.
2. MQA vs GQA: MQA uses 1 KV head (largest cache saving and strongest sharing constraint). GQA uses $g$ KV heads to trade cache size against representational capacity; measure quality.
3. Savings follow the head ratio: KV cache savings are num_query_heads / num_kv_heads. A 64-query-head model with 8 KV heads gets 8x savings; a 32-query-head model with 8 KV heads gets 4x.
4. Uptraining: You can convert MHA checkpoints to GQA by mean-pooling KV heads and continuing training for a small compute budget, then validating quality and serving behavior.
5. Serving Support: Modern GQA-aware kernels such as FlashAttention^[10] and FlashInfer^[11] avoid explicit K/V expansion in the hot path.