Tensor Parallelism

Notes

Tensor parallelism shards the tensors inside one transformer layer across accelerators so a single token’s matmuls are computed cooperatively, with the canonical Megatron-LM scheme from Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” (2019). For a linear layer $Y=XW$, $X\in {\mathbb{R}}^{b\times h}$ and $W\in {\mathbb{R}}^{h\times k}$, column parallelism partitions $W=[W_1,\dots ,W_p]$ along $k$, so device $i$ computes $Y_i=X{W}_i$ and the concatenation $Y=[Y_1,\dots ,Y_p]$ needs no immediate reduction. Row parallelism partitions $W=\left[\begin{array}{c}{W}_1\\ \cdots \\ W_p\end{array}\right]$ along $h$ and $X=[X_1,\dots ,X_p]$, so each device computes $Z_i=X_i{W}_i$ and the layer output is $Y={\sum }_i{Z}_i$, implemented as an all-reduce. In an MLP this pairs naturally: shard the up/gate projections column-wise, keep the expanded activation sharded, then shard the down projection row-wise and all-reduce back to replicated hidden state.

For a Llama-style attention block, TP usually shards Q, K, V, and output projection in the same column-then-row pattern, but grouped-query attention makes the bookkeeping asymmetric. Suppose a 70B model has hidden size $h=8192$, 64 query heads, 8 KV heads, and head dim 128. With $p=8$ tensor-parallel GPUs, each rank owns 8 Q heads and 1 KV head, so the local Q projection has $8192\times 1024$ weights while local K and V projections each have $8192\times 128$ weights; attention is computed locally over the local heads, then the output projection consumes the local $1024$-wide attention result and contributes a partial $8192$-wide residual update that must be all-reduced. vLLM’s PagedAttention paper solved KV-cache paging, not this sharding, but vLLM’s distributed executor uses Megatron-style TP around those per-rank heads; SGLang, TensorRT-LLM, NVIDIA Dynamo, and llm-d expose the same basic “tp size” knob for serving fleets rather than inventing a different algebra.

The cost model is simple enough to do on a whiteboard: TP divides parameter bytes and most matmul FLOPs by $p$, but adds collectives on activation-sized tensors at every block. A ring all-reduce of $n$ bytes across $p$ ranks transfers $2\frac{p-1}{p}n$ bytes per rank, so the communication time is roughly $T_{\text{ar}}\approx 2\frac{p-1}{p}n/B_{\text{link}}+\alpha 2(p-1)$. For decode with batch $b=64$, hidden $h=8192$, FP16 activations, one residual all-reduce after attention output and one after MLP down projection means $n=bh2=64\cdot 8192\cdot 2=1,048,576$ bytes each, or $2$ MiB per layer before ring overhead. At $p=8$, ring traffic per all-reduce is $1.75$ MiB per rank; two collectives over 80 layers move $280$ MiB per generated token per rank, before latency. On an H100 NVLink/NVSwitch domain with effective hundreds of GB/s per collective, this is tolerable; across 200 Gb/s Ethernet, the same payload alone is about $280\text{ MiB}/25\text{ GB/s}\approx 11.2$ ms per token, ignoring software latency.

The compute-side arithmetic explains why engineers still use TP for oversized models. A 70B FP16 model needs about $140$ GB just for weights, so it cannot fit on one 80 GB H100 once KV cache and workspace are included; with $p=2$ the weights become $70$ GB per GPU, barely leaving room, and with $p=4$ they are $35$ GB. On a single 80 GB H100 at $3.35$ TB/s HBM bandwidth, pure weight streaming for one decode token would take $140\text{ GB}/3.35\text{ TB/s}\approx 41.8$ ms if no reuse across batch were possible; with $p=4$ it falls to $35/3.35\approx 10.4$ ms per rank, plus all-reduces. For the earlier $b=64$, $p=4$ case, each all-reduce moves $2(3/4)\cdot 1.0$ MiB $=1.5$ MiB, two per layer over 80 layers gives $240$ MiB per token per rank; at an optimistic $450$ GB/s effective NVLink collective bandwidth, payload time is $0.24/450\approx 0.53$ ms, small next to weight bandwidth and GEMM time.

The technique stops paying when the saved local matmul time is smaller than the collective time and launch/synchronization penalty, which happens first outside one NVLink or NVSwitch island and second at small batches or narrow models. With $b=1$, the all-reduce payload above is only $16$ KiB per collective, but 160 collectives per token for an 80-layer model are latency dominated; at even $20\mu$s each, synchronization alone is $3.2$ ms/token, and over TCP/RDMA Ethernet it is often worse and more jittery. With $p=16$ across two 8-GPU nodes, the ring factor barely changes, $2(15/16)$, but the slowest inter-node hop gates the collective; every layer now waits for the network before the next block can start. TP also becomes awkward when KV heads are fewer than TP ranks: the 8-KV-head example maps cleanly to $p\le 8$, but at $p=16$ some ranks cannot own a whole KV head without head splitting, extra gathers, or uneven work, so TensorRT-LLM and vLLM deployments usually prefer pipeline, data, or expert parallelism beyond the NVLink domain.

FlashAttention from Dao et al. attacks the attention kernel’s HBM traffic by tiling QK softmax V, and EAGLE speculative decoding from Li et al. reduces the number of full-model decode steps accepted per emitted token; neither removes the Megatron all-reduce created by row-parallel projections. PagedAttention in vLLM reduces KV fragmentation and enables high-throughput continuous batching, but its block table is per-rank when TP shards heads. SGLang’s RadixAttention/prefix caching, TensorRT-LLM’s fused collectives and in-flight batching, NVIDIA Dynamo’s disaggregated serving runtime, and llm-d’s Kubernetes-oriented distributed inference stack all have to place TP groups as topology-aware units for the same reason: the model layer graph contains a hard collective edge after attention output and after MLP down projection. A correct fleet plan therefore treats TP as an intra-node memory-and-bandwidth tool, not as a general way to smear one request over arbitrary machines.