Tensor Parallelism Does Not Divide Your KV Cache by N

A 32-billion-parameter model does not fit on one card with enough room left to serve anybody, so you split it across eight. The mental model that comes with that decision is division: eight GPUs, one eighth of everything each, eight times the room. Two of the three things you care about do divide like that. The one you are actually short of stops dividing partway through, and on most current models it stops at exactly eight.

The vocabulary first. Tensor parallelism — Megatron-LM calls it intra-layer model parallelism — cuts every weight matrix in a layer into N pieces and gives one piece to each GPU, so all N run the same layer at the same time on different slices. Attention is cut by head: with 64 query heads and N = 8, each GPU owns 8 of them. The feed-forward network is cut by column and then by row. Because each GPU ends up with a partial sum of the layer's output, the pieces have to be added back together, and that addition is an all-reduce: a collective in which every GPU ends up holding the sum of all N contributions. There are exactly two per decoder layer. And GQA, grouped-query attention, is the now-standard design in which many query heads share one key-value head — Qwen3-32B has 64 query heads and 8 key-value heads — which is what the GQA lesson is about and what makes the KV cache small enough to serve from.

Below is Qwen3-32B on H100 80GB cards, sharded N ways. It is set at --tensor-parallel-size 8, the obvious choice for a single eight-GPU node. Read the hero number — how many 32k-token conversations fit at once — then take the tensor-parallel size to 16.

At --tensor-parallel-size 8 the model holds 60 conversations of 32k tokens at once. Each GPU carries 8.19 GB of weights and 1 key-value head, so 32,768 bytes of cache per token.

Double the GPU count to 16. Capacity goes to 64. Four more conversations — 1.07x — for twice the hardware. Look at why: KV heads per GPU still reads 1, and cache bytes per token per GPU still reads 32,768. Nothing about the cache changed. The only thing the extra eight cards bought was weights room, 8.19 GB per GPU down to 4.18 GB, and weights were never what you were short of. Go to 32 and it is 1.03x.

Now find the cause. Take num_key_value_heads to 64, so every query head has its own key-value head and there is no grouping at all. The same 8-to-16 step now goes from 7 sequences to 16 — 2.29x — and 16 to 32 gives 2.06x. Sharding by head works perfectly right up to the point where there are no more heads to give out, and grouped-query attention is a design whose entire purpose is to leave you with as few as possible. vLLM does not error when it runs out. It replicates: every GPU past the eighth is handed a copy of the same key-value head.

The second half of the bill is in the lower panel. Select PCIe Gen5, 128 GB/s and put the tensor-parallel size back to 8. The prefill chunk goes from 225.7 ms on NVLink to 393.6 ms, and the magenta share of the lane goes from 12.6% to 49.9%. Half the machine's time is now spent adding numbers up. At 32 GPUs on PCIe it is 81.1%, and the chunk takes 268.0 ms — slower than eight NVLink GPUs managed.

What a tensor-parallel layer actually does

Megatron-LM's contribution was noticing that a transformer layer can be split so that only two collectives per layer are needed, and that this can be done "with the insertion of a few communication operations in native PyTorch". The trick is the pairing of two sharding directions.

A column-parallel linear splits the weight matrix by output column: A = [A_1, ..., A_p], so GPU i computes Y_i = X·A_i from the full input and holds a slice of the output. A row-parallel linear splits by input row and takes an input that is already sliced, so GPU i computes a partial sum over its slice and the partial sums must be added. Put a column-parallel layer immediately before a row-parallel one and the slice produced by the first is exactly the slice the second wants — no communication between them.

That is what every decoder layer is built from. In vLLM's Qwen3 implementation the QKV projection is a QKVParallelLinear (column-parallel, sharded by head) and the output projection o_proj is a RowParallelLinear; in the feed-forward network gate_up_proj is a MergedColumnParallelLinear and down_proj is a RowParallelLinear. Both take reduce_results=True, the default. Two row-parallel layers, two all-reduces, per layer, in the forward pass. On a 64-layer model that is 128 all-reduces in the decoder stack on every forward pass — the vocabulary embedding adds one more — whether the pass is prefilling 8,192 tokens or decoding one token for one request.

That last clause is the whole cost structure. The tensor being reduced has shape [tokens, hidden], so its size scales with the number of tokens in the pass. In a ring all-reduce each rank ships 2(N-1)/N of the tensor, which rises from 1.0 at N = 2 towards 2.0 and never falls. The count of collectives, meanwhile, depends only on the layer count. Decoding with a batch of one moves almost no bytes and still pays 128 fixed costs; prefilling a chunk pays almost no fixed cost and moves a great deal of data. Move the fixed cost per all-reduce slider with sequences decoding at 1 and you can watch the first regime; leave it at zero and change the interconnect to watch the second.

Where the sharding stops

Query heads shard perfectly because there are 64 of them and you will never have 64 GPUs in one tensor-parallel group. Key-value heads do not, and vLLM's rule for them is four lines with a comment that gives the whole game away:

# If tensor parallelism is used, we divide the number of KV heads by
# the tensor parallel size. We will replicate the KV heads in the
# case where the number of KV heads is smaller than the tensor
# parallel size so each GPU has at least one KV head.
return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)

max(1, ...). Past tensor_parallel_size == num_key_value_heads the expression is pinned at 1 and every further GPU you add receives a full copy of the same key-value head. The cache is not smaller; it is duplicated. On Qwen3-32B that boundary is at 8, which is also the most common node size, which is why almost nobody meets this until they try to go past one node.

The second thing to be clear about is that capacity is per GPU and not summed. Every GPU in a tensor-parallel group holds a head-shard of every resident sequence — that is what makes it tensor parallelism rather than data parallelism. So the number of sequences the group can serve is the number one GPU's pool can hold, not N times it. The pool grows for two separate reasons as N rises: the weights shrink, freeing bytes, and the cache per token shrinks, so each byte holds more. Only the second of those is worth much, and it is the one that stops.

Third, the set of legal values for N is small and it is model-dependent. vLLM asserts three things at model construction: query heads must divide by N; key-value heads must divide by N, or if there are fewer than N, N must be a multiple of them; and the feed-forward width must divide by N. Drag the tensor-parallel slider onto 3 or 6 and the readouts say refused, with the assertion that fires. This is not a hypothetical: there is a vLLM pull request written by someone who wanted a 70B model on three A100s precisely because it did not fit on two, and it required implementing uneven weight sharding to do it.

A note on what this means for architecture choices you do not control. Multi-head latent attention, the DeepSeek design, sidesteps the floor by compressing the cache into a single latent vector rather than into heads — vLLM's get_num_kv_heads short-circuits to return 1 for it, because "when using MLA during decode it becomes MQA". Multi-query attention, one key-value head, has the floor from the start: set num_key_value_heads to 1 and max(1, 1 // N) is 1 at every N, so the cache per token never shards at all. The seats still climb — 8, 38, 53, 60, 64, 66 across TP 1, 2, 4, 8, 16 and 32 — but every one of those tokens came from weights getting smaller, not cache. The bars turn magenta at TP 8, where the step is 1.13x.

The boundary: when adding GPUs is the wrong purchase

The simulation makes it easy to find the point where more GPUs stop being the answer, and it is worth naming what the alternatives buy instead.

Quantizing the cache attacks the number tensor parallelism has stopped attacking. At the floor you hold 32,768 bytes per token per GPU because you have one key-value head in BF16, the 16-bit brain-float format the weights use; an FP8 (8-bit float) cache halves that on every GPU at once, which is a 2x on capacity where the sixteenth GPU was a 1.07x. That is the trade in what quantization actually changes, and it is the right lever on the right side of the floor.

Data parallelism — running two independent tensor-parallel groups of 8 rather than one group of 16 — does multiply capacity by 2, because the two groups hold different sequences. It costs you the ability to serve any single request faster, and it needs the weights twice. When the sixteenth GPU buys 7% as a tensor-parallel rank and 100% as a second replica, that arithmetic is not close.

Nothing at all is often correct. Put the controls back to the ones the lesson opened with — NVLink 900 GB/s, num_key_value_heads 8, context 32,768 — then set sequences decoding to 256 and read output tokens per second across the sizes: the step from TP 4 to TP 8 is worth 2.28x, and the step from 8 to 16 is worth 1.06x. If the fleet is at TP 8 and saturated, the next machine should be another machine. The interconnect matters less than you would guess here: the same sweep on PCIe Gen5 gives 2.22x and 1.06x, because at a batch of 256 the step is bandwidth-bound on weights and cache, not on the link.

There is one case that runs the other way, and it is the reason nobody should read this lesson as "do not use tensor parallelism". Set sequences decoding to 1 — a single interactive request, latency being the only thing that matters — with the interconnect still on NVLink 900 GB/s and num_key_value_heads still 8. Decode is bandwidth-bound, so splitting the weights over N cards splits the bytes each card must read, and the step time falls from 4.10 ms at TP 8 to 2.60 ms at TP 16 to 1.86 ms at TP 32. Communication climbs from 15.7% of the step to 34.7%, and it still wins. Tensor parallelism is a latency technique that happens to also free memory, and the memory is the part that runs out first. Why splitting bytes helps here at all is a roofline question.

Checking it on a real system

Four numbers, in the order you should look at them.

One: your key-value head count. It is in the model's config.json as num_key_value_heads, and it is the largest tensor-parallel size from which you will get a full memory return. Qwen3-32B and Qwen3-8B both ship 8, which is why the floor and the standard node size are the same number. Read yours rather than assuming it. If your --tensor-parallel-size already equals it, more ranks will not give you more context or more concurrency.

Two: what vLLM says the pool is. The engine prints its own answer at startup, and it is the ground truth that beats any calculation:

INFO ... GPU KV cache size: 1,996,143 tokens, Maximum concurrency for 32,768 tokens per request: 60.92x

One line, not two: vLLM emits both fields from a single logger.info_once. The token count is the real pool after weights and workspace. The second field divides it by max_model_len, so it answers "at full context", not "at the length my requests actually use". The values shown are this lesson's model at --tensor-parallel-size 8, not a capture from a real server.

Run it at your current tensor-parallel size, then at twice that, and compare the two lines. If the token count barely moves, you have found the floor empirically in about four minutes, and no further reasoning is required.

Three: your topology. nvidia-smi topo -m prints a matrix of how each GPU pair is connected. NV# means that many NVLink connections; PIX, PXB, PHB and SYS are all PCIe paths of decreasing quality, and SYS means the traffic crosses the CPU sockets. A tensor-parallel group whose matrix contains SYS is the 128 GB/s case in the simulation or worse. Check this before you tune anything, because it decides whether the magenta half of the lane is 12% or 50%.

Four: the collective's share of the step. Set NCCL_DEBUG=INFO once at startup — NCCL is NVIDIA's Collective Communications Library, the code that actually runs the all-reduce — to confirm which transport was selected — the log names the path per rank pair, and a group you believed was on NVLink silently falling back to shared memory or the network is a common and expensive surprise. For the ongoing measurement, compare median inter-token latency at your tensor-parallel size against the same model at half that size on the same hardware: if halving the group makes per-token latency worse by much less than 2x, communication is eating the difference.

Two configuration notes worth carrying. Tensor parallelism multiplies your exposure to a single slow rank, because 128 all-reduces per forward pass are 128 synchronisation points; one throttling GPU sets the pace for all of them. And the collective implementation is not fixed — vLLM has shipped changes to its default all-reduce path that broke tensor parallelism for particular models, so a version bump that changes throughput without changing your config is worth checking against the release notes before it is worth profiling.

Finally, the interaction that catches people who have already read the rest of this corpus: a larger pool is not by itself more throughput. It raises the ceiling that the scheduler's admission policy runs into, and if --max-num-seqs was already the binding constraint rather than memory, doubling the pool changes nothing you can measure. Check which of the two is binding before buying either.

Next: the architecture that put the floor there, grouped-query attention; the memory it is a floor on, the KV cache; the other way to buy pool space, quantization; and the allocator that hands the pool out block by block, PagedAttention.