Multi-Query & Grouped-Query Attention

The last chapter sized a key-value (KV) cache and showed why decode often waits on memory traffic, not FLOPs. One choice in that formula belongs to the architecture: how many Key and Value heads each token adds to the cache.

Hold eight query heads fixed. Under standard MHA, one new token writes eight K/V pairs. GQA with two KV heads writes two pairs, and MQA writes one. All three still form eight queries; only the amount of persistent K/V state changes.

That smaller state can admit more live requests and reduce bytes read during decode. Sharing also makes several queries use the same K/V representation, so quality can move with the head count. We can calculate the memory win first, then test whether a model and serving stack preserve the quality and latency we need. Paging, scheduling, and quantization address the remaining budget.

In decoder-side self-attention, standard multi-head attention (MHA) gives each query head its own Key and Value projections.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Incremental decode keeps those K/V vectors for every earlier token. As prompts grow and more users share a device, that per-head state can reach hundreds of gigabytes. GQA and MQA shrink the dynamic part by sharing cached K and V across query heads.

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

On each decode step, a model projects the new token and keeps its Key and Value vectors so later steps don't recompute them. Each new token adds one more row to the cache. In standard MHA, every query head owns a K/V row, so the bill grows with head count as well as token count.

Count the stored numbers first

Start with one layer, two attention heads, and four values per head. Count K vectors first, then V vectors, for one token:

Now multiply by eight requests and 2,048 tokens: 16 x 2,048 x 8 = 262,144 numbers for one layer. Across 80 layers, that becomes 20,971,520 numbers before we count bytes. The arithmetic is easy; keeping that state in GPU memory while High Bandwidth Memory (HBM) reads it on every decode step is the serving constraint.

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 of dimension $d_k$, attach a byte count to that picture. The leading 2 stores one Key and one Value tensor. $h$ counts heads, $d_k$ counts values inside each head, and bytes per element comes from the storage format: 2 bytes for FP16 (16-bit floating point), 1 byte for INT8 (8-bit integer), and so on.

Those factors give one token in one layer. Add layers, sequence positions, and requests for a full serving estimate:

$\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)

Put concrete numbers into the formula. Qwen2.5-72B^{[2]Reference 2Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115} has 80 layers, 64 query heads, and a head dimension of 128. Its published architecture uses GQA-8, so the calculation below is a counterfactual MHA cache with the same dimensions:

$\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.)

The first $2$ counts K and V, $80$ counts layers, $4096$ counts tokens per request, and $32$ counts requests. The remaining $64 \times 128$ is the per-head width, and the final $2$ is FP16 bytes per element.

At roughly 344 GB, this cache is more than twice the model's ~144 GB of FP16 weights. For long-context or high-concurrency decode, moving K/V through HBM can dominate incremental attention.^{[3]Reference 3Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}

The same total of 131,072 token positions can arrive as 32 requests at 4K or one request at 128K. The next figure keeps the dimensions and total positions fixed while changing that workload shape, so full MHA still lands near 344 GB.

That is the difference between an MHA cache that doesn't fit on one 80 GB GPU and a GQA-8 cache around 43 GB before weights, runtime buffers, and fragmentation. Neither configuration is a deployment decision by itself: weights, quantization, hardware layout, scheduler policy, and measured traffic still set the budget.

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)

One shared K/V pair

MQA takes the extreme cut. All $h$ query projections remain, but one Key projection and one Value projection serve every query head.^{[4]Reference 4Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150} Storage therefore drops by the full query-head count. The cost is that every query now reads from the same K/V subspace.

In the head-sharing figure above, that's the rightmost column: eight queries, one cached K/V pair.

How much memory MQA saves

Apply the same count to the tiny example: 2 heads, dimension 4, one token. MHA stores 16 numbers. MQA stores one Key and one Value head, so $2 \times 1 \times 4 = 8$ numbers. The reduction is 2x for 2 heads and 64x for 64 heads.

The MHA formula's $h$ becomes $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, so cache per token per layer scales with $d_k$ rather than with $h \cdot d_k$.

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

The byte win has a representational price. Query heads can still ask different questions, but all of them read one shared K/V subspace. Tasks that need several distinct relation or value views may lose quality.

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}

That failure mode gives us a useful test for the next architecture. If latency improves while relation-heavy evals regress, query diversity survived but K/V diversity did not. GQA restores some K/V subspaces without returning to one pair per query head.

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 decode step

Keep projection details out of the first decode sketch. Q has one vector per query head; K and V are single shared sequences. Each query scores the same cached K and blends the shared V, so the cache holds one K sequence and one V sequence rather than $h$ of each.

Why can the two query heads produce different outputs if they share one K/V cache?

Answer

They still have different query vectors, so they put different softmax weights on the same cached keys. Sharing K/V changes stored state, not the fact that each head can ask a different question.

---

Grouped-Query Attention (GQA)

Groups of queries, fewer K/V heads

MQA saves the most memory, but one K/V pair can be too narrow for some workloads. GQA keeps several groups: query heads within a group share one Key and one Value, while different groups keep different K/V pairs. That sits between MQA's one pair for everyone and MHA's one pair per query head.

GQA^{[5]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} uses $g$ KV heads where $1 < g < h$:

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

Each group of $h/g$ queries shares one K and one V. Cached $K$ and $V$ therefore have shape $g \times n \times d_k$, not $h \times n \times d_k$. With $h=8$ and $g=2$, four query heads read each KV pair. Mistral 7B uses the same 32-query, 8-KV shape, so its cache is 4x smaller than MHA's while retaining more than MQA's single shared subspace.^{[6]Reference 6Mistral 7B.https://arxiv.org/abs/2310.06825}

GQA in practice

Published dense layouts make the ratios concrete. They are reference configurations, not a 2026 shortlist. Llama 2 used GQA at 34B and 70B.^{[7]Reference 7Llama 2: Open Foundation and Fine-Tuned Chat Models.https://arxiv.org/abs/2307.09288} Llama 3 kept 8 KV heads across its released dense sizes: 8B has 32 query heads, 70B has 64, and 405B has 128.^{[8]Reference 8The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}

🔬 Research insight: Gemma 2 chose a 2:1 query-to-KV ratio on its 9B model (16 query heads, 8 KV heads) after ablations kept downstream scores close to MHA while improving inference speed.^{[9]Reference 9Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}

Head ratios meet another constraint when you shard a model. Tensor parallelism splits work across accelerators, so a simple head-sharded plan needs query heads to divide evenly by the tensor-parallel (TP) degree.

vLLM divides KV heads while possible. When TP exceeds the KV-head count, it replicates KV ownership so every shard has a KV head.^{[10]Reference 10vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttentionhttps://github.com/vllm-project/vllm} Other runtimes and newer layouts can shard along the sequence dimension too. Treat the helper below as a mental model for divisibility and replication, not as a universal runtime contract.

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

Answer

It's 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 notes

Choosing a ratio

MQA targets serving efficiency: shared KV heads cut cache size and reduce memory traffic during incremental decode.^{[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} GQA keeps several subspaces, but its ratio remains model-specific. Llama 3 kept 8 KV heads even at 405B, so the query-to-KV ratio grew with width.^{[8]Reference 8The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}

There isn't a universal best group count. Model size, quality target, and serving stack decide whether 2:1, 4:1, 8:1, or 16:1 fits the workload.^{[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 team inheriting an MHA checkpoint doesn't have to restart pretraining. It can reduce the K/V heads, initialize the new projections by pooling old ones, and then let the model adapt through additional language-model training. This process is called uptraining.

The pooled checkpoint is an initialization, not an exact preservation of the old attention computation. Plan to validate quality and serving behavior after adaptation.

The conversion has two parts:

1. Mean-pool KV heads: Group existing Key and Value heads into partitions (for example, 64 heads into 8 groups of 8), then average each group's $K$ and $V$ projection matrices. The averages retain signal from the old heads as a useful starting point, but they don't reproduce the old attention computation.
2. Fine-tune (uptrain): Continue the standard next-token objective for a small fraction of the original pretraining budget. The model can then adapt its internal representations to the merged projections.

The original GQA paper^{[5]Reference 5GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} found quality close to MHA after using 5% of the original pretraining compute in T5.1.1 encoder-decoder experiments. That evidence has boundaries: it didn't compare decoder-only models or a same-size GQA model trained from scratch.

A new model can bake its KV-head pattern into the architecture from day one. Uptraining is for an inherited MHA checkpoint where restarting pretraining isn't practical.

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 stored for each position through a learned low-rank latent representation.

How MLA works (high-level)

Instead of caching full per-head Key and Value content vectors for every token, MLA follows this path:

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.^{[11]Reference 11DeepSeek-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.^{[11]Reference 11DeepSeek-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.^{[11]Reference 11DeepSeek-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. DeepSeek-V3 keeps the same cache contract: a 512-wide content latent plus a 64-wide RoPE key.^{[12]Reference 12DeepSeek-V3 Technical Report.https://arxiv.org/abs/2412.19437}

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

Runtime support decides whether the smaller representation helps. GQA slots into grouped-attention implementations and familiar tensor-parallel strategies. MLA changes the cache contract, so it needs a compatible execution path.

Choose GQA when that path is mature; consider MLA when its model-specific kernels and cache handling are ready.

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.^{[11]Reference 11DeepSeek-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 equal throughput, but it raises the number of requests that fit under a fixed cache budget. Using the same 80-layer, 64-query-head, 128-dim-head example and FP16 cache:

These ratios are capacity math, not measured throughput. If KV state is what limits batch size, they approximate the memory-limited ceiling.

Fewer KV heads also mean less HBM traffic per decode step while attention reads the growing cache, which can improve inter-token latency when attention bandwidth binds. That gain isn't automatic TPS: weight streaming, kernels, prefill mix, and the scheduler can dominate. Benchmark the workload before turning a cache ratio into a performance claim.^{[3]Reference 3Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}

Long-context impact

Longer sequences make the same ratio more expensive in absolute bytes. For the same 72B-style dimensions and FP16 cache, per-request state grows linearly with sequence length:

These are derived cache sizes, not claims that a model or runtime supports every listed context. 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.

Head sharing is one part of the end-to-end memory budget, not the budget itself.

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

The serving kernel still has to map the smaller KV-head cache onto the larger query-head set. An index expresses that mapping without copying K/V. Production kernels, including FlashAttention's current flash_attn_with_kvcache path and FlashInfer^{[13]Reference 13FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005}, keep the grouping without materializing repeated K/V.

The index preserves the grouping while the smaller cache stays in place. The next chapter explains the serving-side cache layout; FlashAttention^{[14]Reference 14FlashAttention-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: Model size alone isn't enough for serving. num_kv_heads, context length, and KV-cache dtype often dominate concurrency. GQA-8 raises the memory-limited ceiling by roughly 8x versus same-dimension MHA in this arithmetic example, 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

Once head grouping is correct, many GQA failures come from updating the cache one token at a time. The attention formula can be right while a cursor or position write is wrong.

Trace one decode step in order:

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

The head ratio answers only one question: how much K/V state the architecture creates. The mistakes below happen when that answer gets mistaken for a quality result, a training speedup, or a complete deployment budget.

"Model weights are the same thing as KV cache"

• Symptom: A quantized model fits on paper, 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: Keep weights and KV cache in separate budgets. Add activation buffers, allocator slack, scheduler state, and fragmentation before admitting requests.

"GQA speeds up training"

• Symptom: A training benchmark changes little 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 helps when each new token repeatedly loads cached K/V, not when training computes a full sequence in parallel.

"MQA always destroys quality"

• Symptom: MQA gets rejected even for a small, latency-sensitive model whose task eval shows no visible quality change.

• 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 on your workload. Choose GQA when you need more representational room without returning to MHA's cache cost.

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

• Symptom: A large KV-cache reduction produces only a small latency or cost improvement.

• 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 at realistic batch sizes. Use NVIDIA Nsight Systems or vLLM metrics to identify a memory-bound versus compute-bound workload before attributing a gain to GQA.

"repeat_interleave exploded my memory during GQA training"

• Symptom: A training loop runs out of memory after expanding KV heads with repeat_interleave.

• 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 without materializing the expanded tensor. A custom kernel should broadcast or fuse the grouping instead of creating the expanded tensor.

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

• Symptom: A converted checkpoint loads, but quality drops or tensor-parallel shards disagree during 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. 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 the head-ratio estimate predicts.

• 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.^{[10]Reference 10vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttentionhttps://github.com/vllm-project/vllm}

---

Validate a GQA conversion before rollout

Before rollout, keep the evidence chain short. Preserve published head ordering, convert the checkpoint, and compare cached decoding with full-prefix recomputation on a tiny fixture.

Then compare held-out quality with the original MHA checkpoint and measure physical KV allocation at the intended TP degree and context length. For any throughput or latency claim, record hardware, runtime version, workload, precision, and exact baseline.

Promote only when quality stays inside budget and the measured deployment gains memory headroom or decode throughput without cache drift.