Multi-Query & Grouped-Query Attention

Every active token consumes key-value (KV) cache memory during decode. Model architecture can shrink that cache before serving tricks like paging, scheduling, or quantization are applied.

Picture a transformer serving stack where every attention head keeps its own K/V projection. With five heads, the paperwork pile is small. With sixty-four heads and long contexts, the KV memory grows quickly. One fix is to have groups of heads share one K/V projection instead of giving every picker a personal copy.

That's the idea behind Grouped-Query Attention (GQA) and Multi-Query Attention (MQA). In decoder-side self-attention with standard multi-head attention (MHA), each query head has its own Key and Value projections.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Incremental decode therefore 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 isn't an automatic throughput or quality guarantee.

What do MQA and GQA change compared with standard multi-head attention?

Answer

They keep multiple query heads but reduce the number of cached key/value heads. MQA shares one KV head across all queries, while GQA shares KV heads within groups of query heads.

Why KV cache memory becomes the bottleneck

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 scaling up to billions of parameters, use 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.

For the tiny 2-head example, how much KV cache does one layer need for 8 concurrent requests at 2,048 tokens each?

Answer

Each token stores 16 numbers. The one-layer cache is 16 x 2,048 x 8 = 262,144 numbers. A real model then multiplies again by layers and bytes per number.

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 uses softmax over scaled query-key scores:

$\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$, read the cache size formula term by term 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 the storage format: 2 bytes for FP16 (16-bit floating point), 1 byte for INT8 (8-bit integer), 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]Reference 2Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115} 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]Reference 3Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}

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 model serving, 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.

What is the most common KV-cache sizing mistake?

Answer

Using the per-token-per-layer formula as if it were the full cache. A full serving estimate must multiply by layers, sequence length, batch size, KV heads, head dimension, K/V tensors, and bytes per element.

---

Multi-Query Attention (MQA)

The extreme compression

Back to the K/V-sharing picture: MQA takes the extreme approach. Instead of every query head keeping its own K/V projection, all query heads share one K/V projection. This saves enormous storage, but some lanes might lose details that would have helped their specific work.

MQA^{[4]Reference 4Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150} 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 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.

Your 64-head serving model runs out of memory under long chats. You swap from MHA to MQA. What single term in the KV-cache formula changed, and why does that produce a 64x reduction?

Answer

The KV-head count dropped from 64 to 1. MHA stores K and V for all 64 heads, while MQA stores one shared K head and one shared V head, so the cache shrinks by the same 64x factor.

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]Reference 4Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150[5]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}

Return to the transformer serving stack: MQA gives every query head the same K/V representation. Each query head can still attend to a different aspect of the context (a separate query), but every answer starts from one shared K/V representation. That's efficient, but it can hide a detail that a specialized K/V projection would have preserved.

After an MQA migration, latency is great but relation-heavy evals regress. What bottleneck did you probably introduce?

Answer

All query heads still ask different questions, but they now read from one shared K/V representation. That single shared subspace can bottleneck tasks that need several distinct attention patterns.

A minimal MQA implementation

This basic PyTorch implementation of Multi-Query Attention projects the input into multiple query heads, but uses only a single shared key and value projection across all query heads. The key and value tensors are broadcast across query heads during attention scoring to produce the final output.

Why does the output stay at width 64 even though K and V use only one head?

Answer

The model still keeps four separate query heads and concatenates their attended outputs back into the full model width. Sharing K and V changes cached state, not the final hidden-size contract.

---

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]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} 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 more representational diversity than MQA's single KV pair.

GQA is GPU-side reference sharing: instead of every query head carrying its own K/V cache view (MHA) or all heads sharing one global view (MQA), you organize heads into groups that share one K/V pair. Each group can still ask different queries, but it reads shared cached reference data.

GQA in practice

Production model configurations make the pattern easier to compare:

Beyond memory savings, GQA interacts with tensor-parallel serving. Tensor parallelism splits one model's work across several accelerators. Each runtime has its own sharding rules. For example, vLLM first requires the total query-head count to divide evenly by the tensor-parallel (TP) degree. It divides KV heads across shards while that's possible, then replicates KV heads when TP exceeds the KV-head count so every shard still owns at least one.^{[9]Reference 9vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttentionhttps://github.com/vllm-project/vllm}

A teammate says "32 query heads and 8 KV groups means 8x savings." What is correct, and why?

Answer

It is 4x savings, not 8x. The factor comes from num_query_heads / num_key_value_heads = 32 / 8 = 4, so each KV head is shared by four query heads.

Why can a 64-query-head, 8-KV-head model lose part of its cache saving at TP=16 in vLLM?

Answer

The 64 query heads still divide cleanly across 16 shards, but there are only 8 KV heads. vLLM gives every shard at least one KV head, so each KV head is replicated on two shards. The model still uses GQA, but the physical serving layout duplicates some cached state.

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]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} The correct group count for another model remains an architecture and evaluation choice, not a guarantee inherited from that experiment.

You are serving a large support model and need more concurrency without hurting answer quality too much. Why is GQA often safer than pure MQA?

Answer

GQA keeps several KV subspaces instead of collapsing everything into one shared pair. That preserves more representational diversity while still delivering most of the KV-cache savings.

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]Reference 4Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150[3]Reference 3Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}
• Many modern decoder-only LLMs use grouped KV heads, but the choice is model-specific rather than family-wide. Published examples here include Llama 2 34B/70B, Mistral 7B, Qwen2.5-72B, and Gemma 2 9B.^{[6]Reference 6Llama 2: Open Foundation and Fine-Tuned Chat Models.https://arxiv.org/abs/2307.09288[7]Reference 7Mistral 7B.https://arxiv.org/abs/2310.06825[2]Reference 2Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115[8]Reference 8Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}
• 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]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}

---

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

MHA to GQA conversion: Partition the existing K/V heads into the desired groups, mean-pool each group's projection weights, then continue training so the model adapts to shared K/V representations. The pooled checkpoint is an initialization, not an exact preservation of the original attention computation, so quality and serving behavior both need evaluation.

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 retains information from the original MHA heads as a useful initialization, but it doesn't preserve the original attention computation exactly.
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]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} showed that this recipe reaches quality close to MHA with 5% of the original pretraining compute in its T5.1.1 encoder-decoder experiments. The paper also calls out an important limit: it didn't compare decoder-only models or a same-size GQA model trained from scratch. 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.

How do you initialize GQA from an existing MHA checkpoint?

Answer

Group the old K/V heads and mean-pool their projection weights into fewer KV heads. Then uptrain the converted model so it adapts to the reduced K/V capacity.

---

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 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 doesn't materialize full per-head cached K/V again.^{[10]Reference 10DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434}

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.^{[10]Reference 10DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434}

DeepSeek-V2 reports a 93.3% reduction in deployed KV-cache memory footprint compared with DeepSeek 67B.^{[10]Reference 10DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434} That deployed comparison also includes KV-cache quantization, so don't treat 93.3% as an MLA-only head-count ratio. MLA's architectural payload depends on the chosen latent width $d_c$, positional component, and execution path.

You need long-context serving on a runtime that already supports GQA kernels but not MLA-specific execution. Why is GQA operationally simpler, even if MLA compresses harder on paper?

Answer

GQA keeps a conventional K/V-head contract that existing grouped-attention kernels and sharding layouts support. MLA can compress harder, but the model and runtime must implement its latent-cache attention path correctly.

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.

DeepSeek-V2 documents a 512-dimensional compressed K/V latent and a 64-dimensional decoupled key component. That means MLA stores 576 values per layer and token before storage-precision choices. For a shape comparison, an illustrative GQA layout with 128-dimensional heads stores $2 \times g \times 128$ values, so MLA's 576-value payload sits at the same width as 2.25 GQA heads. This is an architecture-level payload comparison, not a claim that unrelated models will have the same latency or quality.^{[10]Reference 10DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434}

Your stack already ships MLA kernels and cache pressure is still dominant after other optimizations. When does MLA become worth extra complexity?

Answer

MLA becomes worth it when the runtime already supports its fused kernels and the main remaining limiter is KV-cache footprint. In that setting, extra compression can buy more concurrency than GQA alone.

---

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. There is a second effect beyond admission capacity: fewer KV heads mean less HBM traffic per decode step when attention is reading the growing cache, which can improve inter-token latency when attention bandwidth binds. That ITL gain is not automatic TPS: weight streaming, kernels, prefill mix, and the scheduler still dominate many deployments, so treat concurrency ceilings and traffic ratios as capacity math, not a promise of 8× tokens/sec.^{[3]Reference 3Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}

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.

You are targeting 128K context on an 80 GB GPU. Why is GQA-8 alone still not enough, and what else usually joins it?

Answer

GQA-8 cuts the example from about 344 GB to about 43 GB per request, but 43 GB is still too large once weights, batching, fragmentation, and other requests are included. Long-context systems still need paging, quantized KV cache, scheduling, or retrieval.

---

How serving engines handle GQA head expansion

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

Modern GQA-aware kernels, including FlashAttention's current flash_attn_with_kvcache path and FlashInfer^{[11]Reference 11FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005}, handle fewer KV heads without materializing the expanded tensors. This avoids the memory overhead of repeat_interleave and computes grouped attention directly in the fused kernel. The next chapter explains the serving-side cache layout; FlashAttention^{[12]Reference 12FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691} gets its own kernel-level chapter after prefix caching.

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.

Why should a serving engine avoid materializing repeat_interleave for GQA?

Answer

Materializing repeated K/V heads recreates the larger MHA-shaped tensor and wastes memory. Native kernels compute grouped attention directly from the smaller KV-head cache.

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:

In an MHA model, n_kv_heads == n_heads. MQA sets n_kv_heads == 1. GQA sits between them with 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 next-token logits (raw scores before probabilities) should match closely. If they don't, the bug is usually position IDs, mask shape, or cache append order.

Cached decode starts drifting after token 200, while full-prefix recompute still looks correct. Which invariants do you inspect first?

Answer

First, verify the cache shape is (batch, n_kv_heads, max_seq, d_k), not query heads. Then compare cached decode logits against a full-prefix recompute on a tiny prompt, because late-token drift usually comes from wrong position IDs, bad mask shape, or incorrect append order.

---

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 measured latency or cost only improves by a fraction of that.

• Cause: KV cache is one piece of the puzzle. Model weights, feed-forward network (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 isn't 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.

"Architecture-level KV savings always survive tensor parallelism"

• Symptom: Tensor-parallel serving uses more memory than your head-ratio estimate predicted.

• Cause: Physical cache layout is runtime-specific. In vLLM, query heads must divide evenly across TP shards, and KV heads are replicated when TP degree exceeds KV-head count so each shard owns at least one.

• Fix: Check model config and runtime sharding rules before choosing TP degree. Benchmark the physical layout you intend to deploy.^{[9]Reference 9vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttentionhttps://github.com/vllm-project/vllm}

---

Validate a GQA conversion before rollout

Start with a tiny parity fixture: preserve published head ordering, convert the checkpoint, then compare the converted model's cached decoding against its own full-prefix recomputation. Next, compare held-out quality with the original MHA checkpoint and measure physical KV allocation under the intended tensor-parallel degree and target context length. Promote only if quality stays inside budget and realistic concurrency improves memory or decode throughput without cache drift.