DEV Community

jidonglab
jidonglab

Posted on

GQA KV Cache Replication: Why TP=16 Doubles Your Memory

You move Llama-3.1-70B from 8 H100s to 16, expecting roughly double the concurrent sequences. You get about 17% more. The GPUs are not idle, the config has no typos, and nvidia-smi shows memory fully allocated on every rank. The cause is GQA KV cache replication: past a certain tensor parallel size, the KV cache stops sharding and starts duplicating, and every GPU you add carries a full copy of the same key/value tensors.

The threshold is num_key_value_heads. Cross it and your per-GPU KV footprint per token goes flat forever.

TL;DR

  • Grouped-query attention models have few KV heads (8 is the standard for Llama 3.x, Qwen2.5, Mistral, Mixtral). Tensor parallelism shards KV heads across ranks, and a head is indivisible.
  • When tensor_parallel_size > num_key_value_heads, inference engines give each rank a replicated copy of one KV head. Per-GPU bytes-per-token stops shrinking; aggregate KV memory grows linearly with TP.
  • Llama-3.1-70B in bf16 costs 320 KiB/token total, 40 KiB/token/GPU at TP=8 — and still 40 KiB/token/GPU at TP=16. Effective token capacity goes from ~1.39M to ~1.63M. Two times the hardware for ~17% more concurrency.
  • Attention latency also stops improving: TP=16 halves your MLP and QKV projection time but not the per-rank KV read, so attention becomes a hard latency floor while all-reduce cost rises.
  • The fix is pipeline parallelism for the extra dimension (PP=2 × TP=8 shards layers, giving ~3.25M tokens), plus --kv-cache-dtype fp8 to halve bytes-per-token.

What is GQA KV cache replication?

Grouped-query attention (GQA) reduces KV cache size by letting several query heads share one key/value head. Llama-3.1-70B has 64 query heads and 8 KV heads: an 8:1 group ratio, so the KV cache is 8× smaller than it would be under full multi-head attention.

Tensor parallelism splits attention head-wise. Each rank gets num_attention_heads / tp query heads and num_key_value_heads / tp KV heads. That works cleanly while TP divides the KV head count. At TP=8 with 8 KV heads, each rank owns exactly one, and the KV cache is perfectly sharded 8 ways.

At TP=16 there is nothing left to split. A KV head is an atomic 128-dimensional projection; you cannot give half of it to a rank and expect attention to work without an extra collective in the inner loop. So engines take the pragmatic route. vLLM's per-rank head count is shaped like:

# effective KV heads per rank
num_kv_heads = max(1, total_num_kv_heads // tensor_parallel_size)

# TP=8,  8 KV heads -> max(1, 1) = 1   (sharded, 1 unique head per rank)
# TP=16, 8 KV heads -> max(1, 0) = 1   (replicated, 2 ranks per head)
Enter fullscreen mode Exit fullscreen mode

The max(1, ...) is the whole story. Ranks 0 and 8 both hold KV head 0, and both write the same K and V for every token. The configuration is also constrained: TP must either divide the KV head count or be a multiple of it, which is why TP=12 on an 8-KV-head model fails outright while TP=16 silently wastes memory.

Why does GQA KV cache replication happen at all?

Because the alternative is worse. To shard a single KV head across two ranks you would split the head dimension, and each rank would then hold a partial dot product for every query — requiring an all-reduce inside attention, per layer, per decode step. That is a latency disaster at batch sizes where decoding is already communication-sensitive. Replication trades memory for a clean, collective-free attention kernel.

The design is right; the failure is that nobody tells you when you have crossed the line. There is no warning log that says "your KV cache is now 2× larger than necessary." You just see a throughput curve that flattens.

How much memory does GQA KV cache replication actually cost?

Run the arithmetic yourself. Bytes per token for the full model:

2 (K and V) × num_layers × num_kv_heads × head_dim × dtype_bytes
Enter fullscreen mode Exit fullscreen mode

For Llama-3.1-70B in bf16: 2 × 80 × 8 × 128 × 2 = 327,680 B = 320 KiB/token.

def kv_bytes(layers, kv_heads, head_dim, tp, dtype_bytes=2):
    per_rank_heads = max(1, kv_heads // tp)
    replication    = max(1, tp // kv_heads)
    per_rank = 2 * layers * per_rank_heads * head_dim * dtype_bytes
    return per_rank, per_rank * tp, replication

for tp in (2, 4, 8, 16, 32):
    per_rank, total, rep = kv_bytes(80, 8, 128, tp)
    print(f"TP={tp:<3} {per_rank/1024:6.1f} KiB/tok/GPU  "
          f"{total/1024:7.1f} KiB/tok total  x{rep} replication")

# TP=2    160.0 KiB/tok/GPU    320.0 KiB/tok total  x1 replication
# TP=4     80.0 KiB/tok/GPU    320.0 KiB/tok total  x1 replication
# TP=8     40.0 KiB/tok/GPU    320.0 KiB/tok total  x1 replication
# TP=16    40.0 KiB/tok/GPU    640.0 KiB/tok total  x2 replication
# TP=32    40.0 KiB/tok/GPU   1280.0 KiB/tok total  x4 replication
Enter fullscreen mode Exit fullscreen mode

Now translate that into concurrency on 80 GB H100s at the default gpu_memory_utilization=0.9 (~72 GiB usable):

TP=8 — bf16 weights are ~132 GiB, or ~16.5 GiB/GPU. After activations and CUDA graph buffers, roughly 53 GiB is left for KV. At 40 KiB/token that is about 1.39M tokens of cache.

TP=16 — weights drop to ~8.25 GiB/GPU, so about 62 GiB is free. But bytes-per-token is unchanged at 40 KiB. That is about 1.63M tokens.

You doubled the GPU count and bought 17% more KV capacity. The only gain came from the weights getting smaller, not from the cache getting sharded. For a workload serving 32k-token contexts, that is the difference between ~42 and ~49 concurrent sequences — nowhere near the 2× the cluster invoice implies.

Llama-3.1-405B is worse in absolute terms: 126 layers × 8 KV heads × 128 × 2 × 2 = 504 KiB/token, and it is exactly the model people are most tempted to run at TP=16.

Does replication hurt latency too, or just memory?

Both, in an unintuitive split. Going TP=8 → TP=16 halves the work in every dense matmul: QKV projections, the output projection, and the MLP all shard cleanly. Attention does not. Each rank still reads the same 40 KiB/token of KV per decode step it read at TP=8, because it holds a full copy of its head.

Decode attention is memory-bandwidth-bound, so its wall-clock time is set by bytes read. Those bytes did not change. What you get is a model where the dense layers got faster, attention stayed flat, and the fraction of step time spent in attention grew — while the all-reduce after every attention and MLP block now spans 16 ranks instead of 8. On a single NVLink node that is tolerable. Across two nodes over InfiniBand, the added collective latency can eat the entire matmul win, and TP=16 ends up slower per token than TP=8.

That is the specific trap: cross-node tensor parallelism on a GQA model with 8 KV heads is the worst of every dimension at once — replicated cache, flat attention time, doubled collectives, and inter-node hops in the critical path of every layer.

How do I check whether my deployment is replicating KV heads?

One number in config.json decides it:

python - <<'PY'
import json, urllib.request
url = "https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct/raw/main/config.json"
c = json.load(urllib.request.urlopen(url))
print(c["num_key_value_heads"], c["num_hidden_layers"],
      c["hidden_size"] // c["num_attention_heads"])
PY
# 8 80 128
Enter fullscreen mode Exit fullscreen mode

If tensor_parallel_size > num_key_value_heads, you are replicating. It is that simple, and the ratio is your replication factor. Nearly every current open-weight model in the 7B–405B range ships with 8 KV heads, so TP=8 is the ceiling for free KV sharding on most of them. Some models are tighter still — check before assuming.

Cross-check it in the vLLM startup log: it reports the size of the GPU KV cache in tokens/blocks. Halve TP and re-read that line. If the token count barely moves when you double TP, replication is the reason.

What should I use instead of raising tensor parallel size?

Shard the layer dimension, not the head dimension. Pipeline parallelism splits the model by layer, and the KV cache splits with it — no replication, because each stage only stores KV for the layers it owns.

# Replicating: 16 ranks, 8 KV heads -> x2 duplication, ~1.63M tokens of KV
vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 16

# Sharded: TP stays at the KV head count, PP takes the extra dimension
vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 8 \
  --pipeline-parallel-size 2 \
  --kv-cache-dtype fp8
Enter fullscreen mode Exit fullscreen mode

With PP=2 × TP=8, each rank holds 40 layers and one unique KV head: 2 × 40 × 1 × 128 × 2 = 20 KiB/token. Against ~62 GiB of free memory, that is roughly 3.25M tokens — about 2.3× the TP=8 baseline instead of 1.17×. Add fp8 KV cache and it roughly doubles again, to ~6.5M tokens, at the cost of a small quality hit on long-context recall that you should measure on your own eval, not assume.

Pipeline parallelism is not free. It introduces bubbles, so single-stream latency at low batch sizes gets worse, and it works best when you have enough concurrent requests to keep every stage fed. For throughput-oriented serving — which is exactly where KV capacity is the binding constraint — that trade is almost always correct. PP also puts only two point-to-point activation transfers on the wire per token instead of an all-reduce per layer, which makes it the right choice for the cross-node hop specifically.

If you control model selection, note that architectures using latent KV compression (DeepSeek-style MLA) sidestep this entirely: the cache is a single compressed latent per token, not per-head tensors, so it shards on a different axis and does not hit the KV-head wall.

When is KV head replication actually fine?

When you are latency-bound on short contexts and memory is not the constraint. If you serve 2k-token prompts at low concurrency and care about time-per-output-token, TP=16 within a single NVLink domain can still win: the dense-layer speedup is real, and you are never going to fill 62 GiB of KV with short sequences anyway. Replication costs you headroom you were not using.

The failure is specific to KV-bound serving: long contexts, high concurrency, or both. That is when a 2× GPU spend returns 17%.

The short answer

TP=16 doubles your KV cache memory because tensor parallelism shards attention by head, and models with grouped-query attention have only 8 KV heads. Once tensor parallel size exceeds num_key_value_heads, the heads cannot be split further, so inference engines replicate them — every rank stores and reads a full copy. Per-GPU bytes-per-token goes flat (40 KiB/token for Llama-3.1-70B in bf16, at TP=8 and TP=16 alike), aggregate KV memory grows linearly with rank count, and attention latency stops improving even as MLP time halves. Treat num_key_value_heads as the hard ceiling on useful TP, and use pipeline parallelism plus fp8 KV cache for scaling beyond it.

Top comments (0)