Tensor parallelism beyond the KV head count replicates the cache. Derive the per-GPU KV at TP=16 for an 8-KV-head model.

TL;DR: With 8 KV heads, TP=8 already puts one KV head on each GPU; TP=16 cannot split a head further, so each head is stored on 2 GPUs and per-GPU KV stays at $2 \cdot L \cdot 1 \cdot d_h \cdot b$, 63 KiB per token for Llama 3.1 405B, while the cluster holds 2x the logical cache. On 16 H100s with FP8 weights, TP=16 holds 701K tokens; TP=8 with PP=2 on the same GPUs holds 1.40M.

How to approach it

Say how tensor parallelism shards attention before computing anything: query heads, and the KV heads they read, are divided across ranks, and a rank cannot hold a fraction of a head. Then ask the KV head count, the weight precision and whether the 16 GPUs span two 8-GPU nodes, because the last one changes the latency side of the decision as much as the memory side.

A strong answer

Under Megatron-style TP each rank owns $n_q / \text{TP}$ query heads and the KV heads those queries use. With GQA the per-rank KV head count is:

$$n_{kv}^{\text{rank}} = \max\left(1, \frac{n_{kv}}{\text{TP}}\right), \qquad \text{copies of each head} = \max\left(1, \frac{\text{TP}}{n_{kv}}\right)$$

vLLM's attention layers compute exactly this and replicate KV heads when TP exceeds the KV head count. Llama 3.1 405B has 126 layers, 128 query heads, 8 KV heads and head dimension 128:

logical KV per token  = 2 x 126 x 8 x 128 x 2 B   = 516,096 B = 504 KiB
TP=8:  1 KV head/rank = 2 x 126 x 1 x 128 x 2 B   = 64,512 B  = 63 KiB per GPU
TP=16: 1 KV head/rank, each head on 2 ranks       = 64,512 B  = 63 KiB per GPU
cluster bytes per token, TP=16 = 16 x 63 KiB      = 1,008 KiB = 2.0x logical
sanity: query heads per rank at TP=16 = 128/16 = 8, still sharing one KV head

So past TP=8 every extra GPU buys weight capacity and compute but zero KV capacity per token. Each rank still holds every token for its one head. The capacity comparison on 16 H100 SXM 80 GB, FP8 weights (405.85 GB), the 0.92 budget and an assumed 3 GB per-GPU reserve:

budget per GPU   = 80 x 0.92 = 73.6 GB
TP=8 replica     weights 405.85/8  = 50.73 GB, KV 19.87 GB / 64,512 B =   307,978 tok
two TP=8 replicas                  = 2 x 307,978                      =   615,957 tok
TP=16            weights 405.85/16 = 25.37 GB, KV 45.23 GB / 64,512 B =   701,174 tok
TP=8 x PP=2      25.37 GB, 63 layers per GPU,  KV 45.23 GB / 32,256 B = 1,402,348 tok
sanity: PP=2 halves layers per GPU, so bytes per token per GPU halve and tokens double

TP=16 still beats two TP=8 replicas on tokens, which surprises people, because halving the weights per GPU freed 25 GB on every GPU. But half of that freed memory is spent storing a second copy of every key and value. PP=2 frees the same weight memory and keeps one copy, so it holds twice as much.

The latency side points the same way on an HGX box. Two 8-GPU nodes joined by InfiniBand turn every one of the 252 per-token all-reduces (2 per layer) into an inter-node collective, and at small payloads those are latency-bound. PP's cost is different: a single sequence passes through both stages, and the pipeline only fills when enough micro-batches are in flight, which continuous batching provides at high load.

The decision: for an 8-KV-head model on 16 Hopper GPUs, run TP=8 per node with PP=2 (or two replicas, if one fits), not TP=16. The reversal condition: inside one NVLink domain larger than 8 GPUs (a GB200 NVL72 rack), the cross-node penalty disappears, and if the deployment is latency-bound at low concurrency with short contexts, KV capacity is not binding and TP=16's lower weight bytes per GPU can win on time per token despite the duplicate cache.

What interviewers probe next

"Can the engine avoid the duplicate instead?" Yes, by splitting attention differently: data-parallel attention (each rank owns whole sequences and all their KV heads) with TP or expert parallelism only for the MLP. SGLang's enable_dp_attention (opt-in in v0.5.21) is that design.

"Does this happen for Llama 3.3 70B?" The same way: 8 KV heads, so TP=8 and TP=16 both store 40 KiB per token per GPU (2 x 80 x 1 x 128 x 2 B), and only TP up to 8 grows the pool.

"What about an MLA model at TP=8?" Worse in relative terms. The 576-wide latent is shared by all heads and cannot be split by head, so every TP rank holds the full latent for every token, an 8x copy at TP=8. That is why MLA deployments favor data-parallel attention.

How it's scored

Weak: divides the logical 504 KiB by 16 and reports 31.5 KiB per GPU, double-counting the capacity.

Hire: applies $\max(1, n_{kv}/\text{TP})$, gets 63 KiB per GPU at both TP=8 and TP=16, and states that the cluster stores two copies.

Strong Hire: compares TP=16 against TP=8 x PP=2 and two replicas with token counts, explains why TP=16 still beats replicas on capacity, and adds the inter-node all-reduce cost and the DP-attention alternative.

Common mistakes

• Assuming KV per GPU always scales as $1/\text{TP}$. It stops at the KV head count.
• Forgetting the weights side: TP=16 does free memory, so the per-token copy does not make it strictly worse than smaller TP.
• Ignoring that 16 GPUs usually means two nodes and an InfiniBand hop on every all-reduce.
• Treating MLA's latent as splittable by head like GQA's KV heads.