Tensor Parallelism for Serving

TL;DR: Tensor parallelism shards each layer's matrices so that $p$ GPUs each read $1/p$ of the weights per step, and it pays two all-reduces of the hidden state per layer. For Llama 3.1 405B at TP=8 on H200s that is 252 all-reduces per decode step against a 21.1 ms weight floor: tiny 32 KiB messages, so the cost is the fixed latency per call, 2.5 ms at an assumed 10 µs each. In prefill the same all-reduces carry 268 MB each and turn bandwidth-bound.

Where the splits go

Megatron-style TP pairs a column-parallel matrix with a row-parallel one so that only one collective is needed per pair.

Column-parallel matrices split their output dimension, so each GPU produces a slice of the activation with no communication. Attention heads are independent, so a GPU holding heads 0 to 15 of 128 computes attention for those heads alone, reading only its share of the KV cache. The row-parallel matrix that follows splits its input dimension, so each GPU produces a full-width partial sum, and an all-reduce adds them. Two all-reduces per layer, every layer, every step.

The cost model

The payload of each all-reduce is the hidden state: $T \times d \times 2$ bytes in BF16 for $T$ tokens in the step. A ring all-reduce sends $2(p-1)/p$ of the buffer per GPU. Per step:

where $\alpha$ is the fixed cost per call and $\beta$ is one-direction link bandwidth. H100 and H200 NVLink is 900 GB/s bidirectional, so $\beta = 450$ GB/s. Llama 3.1 405B has $L = 126$ and $d = 16{,}384$:

405B, BF16, TP=8 on H200 SXM (141 GB, 4.8 TB/s)
weights/GPU   = 811.7 GB / 8 = 101.5 GB
weight floor  = 101.5 GB / 4.8 TB/s = 21.1 ms per decode step
all-reduces   = 2 x 126 = 252 per step
decode, T = 1:   payload = 16,384 x 2 B = 32 KiB
                 ring bytes/GPU = 2 x 7/8 x 32,768 = 57,344 B -> 0.13 us at 450 GB/s
                 latency term = 252 x 10 us = 2.52 ms (11% of the step)
                 at alpha = 3 us: 0.76 ms (3.5%)
prefill, T = 8,192: payload = 8,192 x 32 KiB = 268 MB
                 ring bytes/GPU = 470 MB -> 1.04 ms each, 263 ms per step
                 compute at peak = 2 x 405.9e9 x 8,192 / (8 x 989e12) = 0.84 s
sanity: decode comm is pure latency, prefill comm is about 31% of ideal compute

The values of $\alpha$ here are assumptions for illustration: roughly what a general collective library pays on tiny messages and what a one-shot custom all-reduce over NVLink can reach. Measure them on the node. Note what TP does and does not divide. The weight floor falls as $1/p$. The latency term stays at $2L\alpha$ whatever $p$ is. So each doubling of TP buys less than the one before, and at high TP the fixed term is what you are optimizing.

Memory is the other half of the story

TP stores one copy of the weights across the group, so it also decides how much KV fits. With vLLM's default gpu_memory_utilization of 0.92 and ignoring activation reserve:

BF16 405B does not fit on eight 80 GB H100s at all: 811.7 GB of weights against 589 GB usable. TP inside one NVLink domain is the precondition for serving it in one node.

The model has 8 KV heads, so TP=8 places exactly one KV head per GPU, 63 KiB per token each. At TP=16, KV heads stop dividing: pairs of GPUs hold the same head, and the cache is replicated twice.

Where it breaks

• Crossing the node. Over InfiniBand, $\alpha$ grows to tens of microseconds, and 252 calls per token ruin decode. Keep TP inside the NVLink domain and use pipeline or data parallelism across nodes.
• Small GEMMs. At TP=8 each GPU's slice of a 70B layer is small enough that GEMV kernels underfill the SMs, so achieved bandwidth falls below the datasheet.
• Stragglers. Every all-reduce is a barrier. One slow GPU, whether from thermal throttling or a noisy neighbor process, sets the pace for all eight.
• Throughput per GPU. For a model that fits at lower TP, two TP=4 replicas usually beat one TP=8 on tokens per GPU, because the same fixed latency term is a smaller share of a longer TP=4 step.

What interviewers listen for

• Two all-reduces per layer, and where they sit (after the O projection and after the down projection).
• 450 GB/s per direction, not 900, in the bandwidth term.
• The decode all-reduce recognized as latency-bound from its size, 32 KiB at $d = 16{,}384$.
• KV heads limiting TP before attention heads do.
• TP justified by latency or memory, never assumed to be free throughput.

Say it in 30 seconds

TP splits QKV and the MLP's first matrices by columns and the output matrices by rows, so each layer needs two all-reduces of the hidden state. That divides the weight bytes per GPU by $p$: Llama 405B at TP=8 on H200s reads 101.5 GB per step, a 21.1 ms floor. It does not divide the 252 per-step all-reduce latencies, about 2.5 ms at 10 µs each, which is why doublings give diminishing returns. Prefill all-reduces are 268 MB and bandwidth-bound. Keep TP inside NVLink and stop at the KV head count.