Grouped Query Attention: Why 8 KV Heads Decide Your Batch Size

Every open-weight model you serve — Llama 3, Mistral, Qwen, Gemma — has 8 key/value heads. Not 32. Not 64. Eight. That number is not a coincidence, and it is not a tuning artifact. It is the single hyperparameter that decides how many concurrent requests fit on your GPU, and it was chosen to divide evenly into the tensor-parallel degrees people actually deploy.

Grouped query attention (GQA) is why a 70B model can hold 500k cached tokens on four H100s instead of 60k. It is also why your fine-tuned model silently loses its decode speedup if the attention kernel is wrong, and why --tensor-parallel-size 16 quietly stops saving you memory.

TL;DR

• GQA shares one K/V head across a group of query heads. With 64 query heads and 8 KV heads, the group size is 8, and the KV cache shrinks by exactly 8x versus multi-head attention (MHA).
• KV bytes per token = 2 × layers × kv_heads × head_dim × dtype_bytes. For Llama 3 70B in fp16 that's 320 KiB/token — versus 2.5 MiB/token if it used MHA. That ratio is your batch size multiplier.
• Decode is memory-bandwidth-bound, not FLOP-bound. GQA raises arithmetic intensity by the group size because one loaded K/V tile serves 8 query heads.
• The tensor-parallel trap: when tp_size > num_kv_heads, engines like vLLM replicate KV heads across ranks. TP=16 on an 8-KV-head model gives you zero additional KV memory savings per GPU.
• The repeat_kv trap: reference implementations expand KV back to full head count before the matmul. Cache storage still shrinks, but the bandwidth win at the kernel level disappears unless you use a GQA-native kernel.

What is grouped query attention, exactly?

Grouped query attention is the interpolation between multi-head attention (every query head has its own K and V projection) and multi-query attention (all query heads share one K and one V). You pick num_kv_heads somewhere in between, and each KV head serves a contiguous group of num_attention_heads / num_kv_heads query heads.

The query projection is untouched. Only k_proj and v_proj shrink. In Llama 3 8B, q_proj maps 4096 → 4096 (32 heads × 128), while k_proj and v_proj map 4096 → 1024 (8 heads × 128). The attention math is identical apart from which K/V slice each query head reads.

Why does this work at all? Because during autoregressive decode, the expensive part is not computing attention scores — it is reading the entire KV cache from HBM for every single token you generate. Query heads are cheap; distinct KV heads are not.

Multi-query attention (Shazeer, 2019) took this to the extreme with one KV head and paid for it in quality and training instability. GQA (Ainslie et al., 2023) found the knee: a handful of KV heads recovers nearly all MHA quality while keeping most of the MQA memory win.

Why does GQA decide your batch size?

Because KV cache is the only thing in your memory budget that scales with both sequence length and concurrency. Weights are fixed. Activations are small during decode. The cache is everything else.

def kv_bytes_per_token(layers, kv_heads, head_dim, dtype_bytes=2):
    # factor 2 = one tensor for K, one for V
    return 2 * layers * kv_heads * head_dim * dtype_bytes

# Llama 3 70B: 80 layers, 64 query heads, 8 KV heads, head_dim 128
gqa = kv_bytes_per_token(80, 8, 128)        # 327,680  = 320 KiB/token
mha = kv_bytes_per_token(80, 64, 128)       # 2,621,440 = 2.5 MiB/token

def max_concurrency(free_bytes, per_token, ctx_len):
    return free_bytes // (per_token * ctx_len)

free = int(160 * 1024**3)                    # 4xH100 (320GiB) - 140GB fp16 weights - overhead
print(max_concurrency(free, gqa, 8192))      # 64 concurrent 8k-token sequences
print(max_concurrency(free, mha, 8192))      # 8

Sixty-four versus eight. Same GPUs, same weights, same context length. That 8x is your throughput ceiling, because on a saturated server, throughput is roughly (concurrent sequences) × (decode rate), and the decode rate barely moves until you are compute-bound.

Note what does not appear in that formula: num_attention_heads. You can double the query heads and the cache does not grow by a byte. This is why modern configs look lopsided.

Model	Layers	Q heads	KV heads	Group	fp16 KV/token
Llama 3 8B	32	32	8	4	128 KiB
Mistral 7B	32	32	8	4	128 KiB
Llama 3 70B	80	64	8	8	320 KiB

Run the 8B number against a 128k context: 128 KiB × 131,072 tokens = 16 GiB of KV cache for one request. On an 80GB card holding 16GB of weights, that's four concurrent long-context users. With MHA it would have been one — with no room for the weights.

Why does GQA speed up decode, not just save memory?

Because decode attention has terrible arithmetic intensity, and GQA multiplies it by the group size.

Generating one token with a batch of 1 requires reading the entire KV cache — potentially gigabytes — to do a few million FLOPs. You are nowhere near the GPU's compute roofline; you are pinned to HBM bandwidth. The ratio of FLOPs to bytes moved is what determines whether you're wasting the chip.

With MHA, each query head loads its own K/V tile: one tile read, one head's worth of math. With a group size of 8, the kernel loads a K/V tile once and reuses it for 8 query heads. Same bytes, 8x the math. That's a direct 8x improvement in attention's arithmetic intensity, which is exactly the axis that matters during decode.

This is also why GQA pairs so well with FlashAttention-style kernels: the K/V tile stays resident in SRAM while the grouped queries stream past it.

What breaks when num_kv_heads is smaller than your tensor-parallel degree?

Tensor parallelism shards attention by head. With 8 KV heads and TP=8, each rank owns exactly one KV head and stores 1/8 of the cache. Clean.

Now set TP=16. There is no way to give each rank a fraction of a head, so inference engines replicate KV heads: two ranks hold identical copies of the same KV head. vLLM does this deliberately rather than failing. The consequence is that doubling your GPU count past num_kv_heads buys you weight-sharding and more aggregate bandwidth, but zero additional KV capacity per GPU. Your per-GPU cache footprint stops shrinking.

This is the actual reason every major open model landed on 8. It divides evenly into 1, 2, 4, and 8 — the TP degrees that map onto real nodes. If you are training a custom model and pick num_kv_heads=6 because it looked fine on a quality ablation, you have made it undeployable at TP=4 without replication waste.

Check it before you commit to a topology:

# Llama 3 70B, 4-way TP: 8 KV heads / 4 ranks = 2 KV heads per rank. Good.
vllm serve meta-llama/Meta-Llama-3-70B-Instruct \
  --tensor-parallel-size 4 \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.92

# TP=16 on the same model: KV heads get replicated 2x.
# Per-GPU KV footprint is unchanged vs TP=8.

Why doesn't GQA speed up my fine-tuned model?

Almost always because the attention path calls repeat_kv and materializes the expanded tensor. The reference PyTorch implementation in HuggingFace Transformers does this:

def repeat_kv(hidden_states, n_rep):
    b, num_kv_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        b, num_kv_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(b, num_kv_heads * n_rep, slen, head_dim)

That .reshape() after .expand() cannot stay a view — it allocates. You just wrote a full MHA-sized K and V tensor to HBM and then read it back. Cache storage is still 8x smaller (the persistent cache holds 8 heads), but the per-step bandwidth advantage is gone, and you burn transient memory proportional to n_rep.

Fix it by using a kernel that understands grouping natively. PyTorch's scaled_dot_product_attention exposes an enable_gqa flag in recent versions; FlashAttention and the vLLM/SGLang paged kernels handle mismatched Q/KV head counts directly.

# Instead of expanding K/V, let the kernel broadcast over the group.
out = F.scaled_dot_product_attention(
    q,          # [b, 32, s, 128]
    k, v,       # [b,  8, s, 128]  <- left as-is
    is_causal=True,
    enable_gqa=True,
)

If you wrote custom attention for a research fork, this is the first thing to audit. It looks correct, passes your tests, and silently costs you the entire point of GQA.

One related shape trap: LoRA on a GQA model. k_proj and v_proj have output dimension kv_heads × head_dim, which is 4x or 8x narrower than q_proj. Custom merge or quantization code that assumes square attention projections will either throw or, worse, broadcast wrong.

How do you convert an MHA checkpoint to GQA?

Mean-pool, then uptrain. The GQA paper's construction takes the existing MHA key and value projections, averages the heads within each target group into a single head, and continues pretraining for a small fraction of the original compute budget. Averaging beats picking one head from the group or random init, because the pooled head starts close to the group's shared subspace.

Practically, this is worth knowing for two reasons. First, it means you can retrofit older MHA checkpoints rather than retraining. Second, it explains the quality curve: GQA with a moderate number of groups lands very close to MHA, while collapsing all the way to one KV head (MQA) shows a clearer degradation, especially on tasks that need many distinct retrieval patterns per layer. More KV heads means more independent "lookup channels" per layer; the shared query heads then differ only in what they ask, not where they look.

When should you not use GQA?

When you can afford something better. GQA is a blunt instrument — it removes KV heads outright. DeepSeek's Multi-head Latent Attention takes the other path: keep the head diversity but store a low-rank latent that gets projected back up per head, so the cached tensor is smaller than even aggressive GQA while preserving distinct per-head behavior. It costs extra compute at decode and a more complex kernel, but the cache-per-token numbers are dramatically better for long-context workloads.

If you are serving existing open weights, this is not a choice you get to make — the architecture is baked in. Your levers are the ones downstream: TP degree that divides num_kv_heads, a GQA-native kernel, paged allocation, and cache dtype.

So, why do 8 KV heads decide your batch size?

Because in grouped query attention the KV cache — the only memory term that grows with both concurrency and context — scales with num_kv_heads, not with query heads. Cutting 64 KV heads to 8 cuts Llama 3 70B's cache from 2.5 MiB to 320 KiB per token, which turns roughly 8 concurrent 8k-context requests into roughly 64 on the same four GPUs, and simultaneously raises decode arithmetic intensity 8x by letting one loaded K/V tile serve eight query heads. Eight specifically, because it divides evenly into every tensor-parallel degree people deploy — past that point engines replicate KV heads and the per-GPU savings stop. Verify two things in production: that tensor_parallel_size divides num_kv_heads, and that your attention kernel broadcasts over groups instead of calling repeat_kv. Either mistake gives back the entire benefit while your config file still says num_key_value_heads: 8.