Tensor parallelism

You launch Llama-3-70B with --tensor-parallel-size 4 on an 8×H100 node and the engine reports a KV pool that holds roughly 567k tokens. You bump to --tensor-parallel-size 8 and it reports about 1.56M — a 2.76× jump from a 2× change in GPU count, which feels like a bargain. So you wire up a second node, set --tensor-parallel-size 16, and the number lands at about 1.78M. You doubled the hardware and bought 14% more cache. Every one of those figures is derived in §4 from published shapes, but the shape of the curve is the point: tensor parallelism buys three different things — weight room, FLOP/s, and KV capacity — and they stop arriving at three different times.

The first one to stop is KV capacity, and it stops at a number written in the checkpoint: $h_{kv}$, the number of key/value heads. Llama-3-70B has eight. Past TP=8 the engines replicate KV heads instead of splitting them, because there is nothing left to split. Meanwhile the per-step communication bill keeps growing with every rank you add, so beyond that point you are paying more collectives for a cache that no longer grows.

The second problem is that a naive sharding of a transformer block wants a collective after every matmul — four per block, 320 for an 80-layer model, on a decode step whose entire bandwidth-bound floor is about 5.2 ms at TP=8. Getting that to two per block is the single structural decision in this chapter, and it comes out of two lines of linear algebra.

Hold one object in your head: the residual stream. It is a [num_tokens, d] tensor that every rank must agree on, byte for byte, at the boundary between blocks. Everything between two consecutive points on that stream — the QKV projection, attention itself, the out-projection, the MLP — can be cut into $p$ independent pieces that never talk to each other. The all-reduce is the price of putting the stream back together.

So the design question is not "how do I shard a matmul". It is "how few times per block must the ranks agree?" The answer for a standard pre-norm transformer block is two: once at the end of attention, once at the end of the MLP. And the reason is that each of those two sub-blocks is a matmul, a nonlinearity, and another matmul — and there is exactly one way to split that sandwich so the nonlinearity in the middle does not force a collective.

Take the MLP, because its shapes make the argument loudest. Write it as $Y = \sigma(XA)B$ with $X \in \mathbb{R}^{n \times d}$ the token batch, $A \in \mathbb{R}^{d \times f}$ the up-projection, $B \in \mathbb{R}^{f \times d}$ the down-projection, $\sigma$ an elementwise nonlinearity, $n$ the number of tokens in the pass, $d$ the model width, $f$ the FFN width, and $p$ the tensor-parallel size. For Llama-3-70B, $d = 8192$ and $f = 28672$, so the intermediate is 3.5× wider than the residual stream. Remember that ratio; it decides everything below.

Split $A$ by columns. $A = [A_1 \,|\, \cdots \,|\, A_p]$ with each $A_i \in \mathbb{R}^{d \times f/p}$. Rank $i$ holds a replicated $X$ and computes $Z_i = XA_i \in \mathbb{R}^{n \times f/p}$. No communication: $X$ is already everywhere.

The nonlinearity survives it. $\sigma$ is elementwise, so $\sigma(Z)$ restricted to columns $[i\cdot f/p, (i{+}1)\cdot f/p)$ is exactly $\sigma(Z_i)$. Rank $i$ can apply the activation to its own slice and be correct. This is the whole trick, and it is why the ordering is not arbitrary.

Split $B$ by rows. $B = [B_1; \cdots; B_p]$ with each $B_i \in \mathbb{R}^{f/p \times d}$. The matching rows of $B$ are exactly the ones that multiply rank $i$'s columns of $\sigma(Z)$, so

Each rank produces a full-width $[n, d]$ tensor that is a partial sum. Adding them elementwise across ranks gives the answer. That addition is one all-reduce, and it is the only one in the module.

Why every other pairing is worse

There are four ways to pair the two splits, and the cost differences are not subtle. Two independent facts drive them: (1) $\sigma$ commutes with a column split but not with a row split, because $\sigma(\sum_i Z_i) \neq \sum_i \sigma(Z_i)$ — so a row-parallel first matmul forces a collective before the activation; (2) that forced collective lands on the $[n, f]$ intermediate, which is 3.5× the bytes of the $[n, d]$ residual.

Column→row is the only pairing whose collective never touches the wide intermediate. That is the design, and it is why ColumnParallelLinear and RowParallelLinear come in pairs in every model file in both engines. Attention has the same sandwich shape — a projection, a per-head operation, a projection — so it gets the same treatment, which is why the count is two all-reduces per block and $2L = 160$ per forward pass for Llama-3-70B.

Attention is column→row too, but the column index has a physical meaning. The QKV projection's output is $(h + 2h_{kv})d_h$ wide, and $\mathrm{softmax}(QK^\top/\sqrt{d_h})V$ is block-diagonal in the head axis: head $j$'s output depends on head $j$'s $Q$, $K$, $V$ and nothing else. So the head axis is the column axis, attention is the "nonlinearity" in the middle of the sandwich, and the out-projection — whose input is the concatenation of head outputs — is row-parallel along exactly the same boundary.

Both engines write the head division the same way, in the model file, before the layer is even constructed:

        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

        tp_size = get_parallel().tp_size
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

Character-for-character the same arithmetic, comments included — SGLang's linear and model layers began as a fork of vLLM's and the head-division logic has not diverged at these SHAs. The load-bearing token is max(1, ...). Query heads divide exactly, forever, until $p = h$. KV heads divide only until $p = h_{kv}$, and then every additional rank gets a duplicate. vLLM's QKV layer names the duplication factor explicitly:

        # Divide the weight matrix along the last dimension.
        tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
        self.num_heads = divide(self.total_num_heads, tp_size)
        if tp_size >= self.total_num_kv_heads:
            self.num_kv_heads = 1
            self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
        else:
            self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
            self.num_kv_head_replicas = 1

Put the two effects in one table. Assume an 80 GiB H100, 4 GiB reserved for activations, workspace and CUDA graphs, and the rest given to the KV pool. Llama-3-70B is 131 GiB of bf16 weights and 320 KiB of KV per token at $h_{kv}=8$ (§2.1), so the per-rank cell is $320/\min(\mathrm{TP}, 8)$ KiB.

This is the cluster capacity law that §3.5 derives: $N_{\text{GQA}} = P \cdot \min(\mathrm{TP}, h_{kv}) / c$. The table is that formula with the weight term folded into $P$, which is why the TP=2→4 gain is a wild 8.24× (the pool itself quadruples as the weights vacate) and the TP=8→16 gain is 1.14× (only the weight term is still moving). §3.5 also establishes the counterpart you must not confuse with this: MLA caches are replicated on every rank at all TP, verified in both engines, so an MLA model gets none of this curve and its cache advantage over GQA-8 is spent by TP ≈ 4.7.

Four things stay whole, and each for a different reason.

The vocab-parallel passage, and why 128k matters

Llama-3's vocabulary is 128,256 tokens. At $d = 8192$ the LM head is $128256 \times 8192 = 1.05$ G parameters — 2.10 GB in bf16, and the untied input embedding is another 2.10 GB. Replicating those on eight ranks would cost 33.6 GB of aggregate HBM instead of 4.2, but the sharper cost is bandwidth: a decode step reads every weight it touches, so a replicated head would add $2.10 - 0.26 = 1.84$ GB of extra reads per rank per step, which at 3.35 TB/s is 549 µs — a 10.5% tax on a 5.25 ms step floor (derived; the floor comes from §0.4). So it is sharded.

The input side is elegant. Each rank masks out ids it does not own, embeds the rest, zeroes the masked rows, and all-reduces:

    def forward(self, input_):
        if self.tp_size > 1:
            # Build the mask.
            masked_input, input_mask = get_masked_input_and_mask(
                input_,
                self.shard_indices.org_vocab_start_index,
                self.shard_indices.org_vocab_end_index,
                self.shard_indices.num_org_vocab_padding,
                self.shard_indices.added_vocab_start_index,
                self.shard_indices.added_vocab_end_index,
            )
        else:
            masked_input = input_
        # Get the embeddings.
        output_parallel = self.quant_method.embedding(self, masked_input.long())
        # Mask the output embedding.
        if self.tp_size > 1:
            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
            # Reduce across all the model parallel GPUs.
            return tensor_model_parallel_all_reduce(output_parallel)
        return output_parallel

SGLang's is the same shape, with the extra branch that sends the reduce to the attention-TP subgroup when DP-attention is on — a mechanism owned by §5.3:

    def forward(self, input_):
        # Surface a bad token id (>= vocab_size, or a negative / unmasked sentinel) as a
        # located async assert instead of a silent OOB embedding gather (tp=1 does not mask).
        maybe_detect_oob(
            input_, 0, self.num_embeddings, "VocabParallelEmbedding input id"
        )
        output_parallel = self._embed_local_shard(input_)
        if self.tp_size > 1 and not get_attn_tp_context().input_scattered:
            if self.use_attn_tp_group:
                output_parallel = attn_tp_all_reduce(output_parallel)
            else:
                # Reduce across all the model parallel GPUs.
                output_parallel = tensor_model_parallel_all_reduce(output_parallel)
        return output_parallel

That all-reduce is on a $[n, d]$ tensor — 16 KiB per token, same as every other one, and it happens once per forward pass rather than 160 times. The output side is where the vocabulary hurts. vLLM's logits path:

    def _get_logits(
        self,
        hidden_states: torch.Tensor,
        lm_head: VocabParallelEmbedding,
        embedding_bias: torch.Tensor | None,
    ) -> torch.Tensor | None:
        # Get the logits for the next tokens.
        logits = self._apply_head(lm_head, hidden_states, embedding_bias)

        # Gather logits for TP
        if lm_head.tp_size > 1:
            logits = self._gather_logits(logits)

        # Remove paddings in vocab (if any).
        if logits is not None:
            logits = logits[..., : self.org_vocab_size]
        return logits

Price it. Full logits for one token in bf16 are $128256 \times 2 = 250.5$ KiB. A ring all-gather moves $(p-1)/p$ of the gathered size per rank, so at TP=8 that is 219 KiB per token — versus 28 KiB for one of the block all-reduces. The single logits all-gather costs 7.8× what an entire transformer block's attention all-reduce costs, and it is 4.9% of the model's whole per-token TP traffic for one layer out of eighty.

vLLM has an escape hatch for the common case where you only need the argmax. Each rank takes its local max, and only (value, index) pairs cross the wire:

    ) -> torch.Tensor:
        """Vocab-parallel argmax without all-gathering full logits.

        Each TP rank computes local argmax, then only the (value, index) pairs
        are gathered and reduced. Communication: O(batch * 2 * tp_size) vs
        O(batch * vocab_size).
        """

The gathered tensor is $2p$ fp32 values = 64 bytes, of which 56 cross the wire per rank. That is a 4,000× reduction in logits communication (derived: 224,448 / 56). It is exposed through LocalArgmaxMixin, which LlamaForCausalLM inherits (vllm/model_executor/models/llama.py:L446-L455) and which is used for draft heads in speculative decoding, where greedy argmax is the whole job.

Both engines pad the vocabulary before sharding — DEFAULT_VOCAB_PADDING_SIZE = 64 in both (vllm/model_executor/layers/vocab_parallel_embedding.py:L32, python/sglang/srt/layers/vocab_parallel_embedding.py:L48) — so that the shard boundary lands on a tensor-core-friendly multiple and every rank gets the same shape. 128,256 is already $64 \times 2004$, so Llama-3 needs no padding; a model with an awkward vocab gets rows of garbage that logits[..., : self.org_vocab_size] then slices off.

A ring all-reduce over $p$ ranks does a reduce-scatter followed by an all-gather, each moving $(p-1)/p$ of the tensor per rank. So for a tensor of $S$ bytes, each rank moves

For a transformer forward pass with $L$ layers, two all-reduces each, $n$ tokens, width $d$, $b$ bytes per element, the per-rank bill is

Llama-3-70B, $L = 80$, $d = 8192$, $b = 2$: one token's residual slice is 16 KiB, and at TP=8 that is $160 \times 1.75 \times 16\,\text{KiB} = 4.375$ MiB per rank per token. NVIDIA rates fourth-generation NVLink on H100 SXM at 900 GB/s per GPU, bidirectional, with an NVSwitch fabric giving every GPU in an 8-way HGX node that rate to any other (NVIDIA NVLink). Against that, use 450 GB/s in one direction for this sent-byte convention: 4.375 MiB takes 10.19 µs under the ideal link model.

Compare it to the decode step's floor. From §0.4, a decode step cannot beat weight bytes over HBM bandwidth: $131\,\text{GiB} / 8 / 3.35\,\text{TB/s} = 5.25$ ms. (That charges all 131 GiB. §5.2 quotes 5.17 ms for the same step because it drops the 2.1 GiB input embedding, which is gathered a row at a time rather than streamed; the 1.5% moves nothing below.) So the bandwidth term of TP communication is 0.194% of the floor at batch 1 — nothing. Rearranged, the comm fraction has a pleasingly simple form:

where $P$ is the parameter count, $\beta$ the HBM bandwidth and $B_{\text{net}}$ the interconnect bandwidth. Note what is not in it: the bytes-per-element $b$ cancels, because it multiplies the wire volume and the weight stream equally, and so does the $1/p$ that both terms carry. What is left is linear in batch size and linear in $p-1$. That single expression is the whole scaling story.

Two reasons the scaling is sublinear

Bandwidth. Per-rank compute falls as $1/p$ but per-rank communication rises toward a ceiling: $2(p-1)/p \to 2$. At $p=2$ you move $1.0S$; at $p=8$, $1.75S$; at $p=32$, $1.94S$. So the ratio of comm to compute grows roughly as $p$, exactly as the formula says. At batch 256 the comm bill is 49.8% of the floor at TP=8 and 220% at TP=32 — past which the collectives alone exceed the time the arithmetic was supposed to take.

Latency. At batch 1 each all-reduce is 16 KiB, far too small to reach peak bandwidth; it is a latency event, not a transfer. The step budget is 5.25 ms across 160 collectives, so 32.8 µs per collective before communication consumes the whole floor. If a small NVLink all-reduce costs $\lambda$, the tax is $160\lambda / 5249\,\mu\text{s}$: at $\lambda = 5\,\mu\text{s}$ that is 15%, at $\lambda = 10\,\mu\text{s}$ it is 30%. Neither number is measured here — the point is that $\lambda$ is a constant that TP does not amortise, so scaling efficiency decays even with zero bandwidth cost. Reducing $\lambda$ is precisely what custom all-reduce kernels exist for, and that is §5.4's subject. Measuring it is Lab 07.

Where TP stops paying: the node boundary

Define an interconnect intensity: FLOPs of compute per byte of TP traffic, per rank. Dividing the pass's $2P/p$ FLOPs by $4L\frac{p-1}{p}db$ bytes per token gives

and the machine's own ratio is $\pi / B_{\text{net}}$. For Llama-3-70B at TP=8, $I_{\text{net}} = 3{,}815$ FLOP/byte. H100 NVLink gives $989.4\,\text{TFLOP/s} / 450\,\text{GB/s} = 2{,}199$, so this ideal compute/link ratio is about 1.7; the break-even is at $p \approx 13$, comfortably past the 8-GPU NVLink domain. Cross a node boundary onto InfiniBand NDR at 400 Gb/s = 50 GB/s (NVIDIA Quantum-2) and the machine ratio becomes 19,788 — break-even at $p \approx 2.35$. This example favors keeping TP inside the assumed eight-GPU NVLink domain; it is not a prohibition on multi-node TP. Network topology, achieved bandwidth, workload and model capacity can change that choice. Past the node you use pipeline parallelism (§5.2) or data parallelism (§5.3), which move a $[n,d]$ activation once per stage boundary instead of 160 times per step.

The declaration is a constructor argument, and the collective is issued inside the layer. That is the entire API surface in both engines. vLLM's column layer divides its output:

        # Divide the weight matrix along the last dimension.
        if disable_tp:
            self.tp_rank, self.tp_size = 0, 1
        else:
            self.tp_rank = (
                tp_rank if tp_rank is not None else get_tensor_model_parallel_rank()
            )
            self.tp_size = (
                tp_size
                if tp_size is not None
                else get_tensor_model_parallel_world_size()
            )
        self.input_size_per_partition = input_size
        self.output_size_per_partition = divide(output_size, self.tp_size)
        self.output_partition_sizes = [self.output_size_per_partition]
        # If QKV or MergedColumn, use output size of each partition.
        if hasattr(self, "output_sizes"):
            self.output_partition_sizes = [
                divide(output_size, self.tp_size) for output_size in self.output_sizes
            ]

and the row layer divides its input and, at the end of forward, issues the only collective in the module:

        # Matrix multiply.
        # Only fuse bias add into GEMM for rank 0 (this ensures that
        # bias will not get added more than once in TP>1 case)
        bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
        output_parallel = self.quant_method.apply(self, input_parallel, bias_)

        if self.reduce_results and self.tp_size > 1:
            output = tensor_model_parallel_all_reduce(output_parallel)
        else:
            output = output_parallel

which is a two-line dispatch to the TP process group (vllm/distributed/communication_op.py:L12-L14):

def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
    """All-reduce the input tensor across model parallel group."""
    return get_tp_group().all_reduce(input_)

SGLang's RowParallelLinear.forward has the same skeleton and three extra branches, which is where the engines actually diverge:

        # skip_all_reduce: explicit call-site override. Also honor
        # ForwardFlags (fuse_mlp_allreduce / mlp_reduce_scatter) published by
        # the decoder — callers should not thread those flags into modules.
        if (
            ((self.reduce_results and self.tp_size > 1) or self.use_decode_attn_tp)
            and not skip_all_reduce
            and not should_skip_mlp_all_reduce()
        ):
            if self.use_dp_attention_reduce:
                output = get_parallel().attn_tp_group.all_reduce(output_parallel)
            else:
                quantize_communications = (
                    (
                        not forward_batch.forward_mode.is_decode_or_idle()
                        and get_exec().comm.enable_quant_communications
                    )
                    if forward_batch is not None
                    else False
                )
                if quantize_communications:
                    output = tensor_model_parallel_quant_all_reduce(output_parallel)
                else:
                    output = tensor_model_parallel_all_reduce(output_parallel)

Three differences worth naming. First, use_dp_attention_reduce redirects the reduce to a subgroup — the attention-TP group rather than the global TP group — which is the hook for DP-attention (§5.3). Second, quantize_communications can send the partial sums in a narrower dtype during prefill, trading numerics for wire time, gated off for decode where the message is too small for bandwidth to matter — exactly the regime split Figure 3 predicts. Third, should_skip_mlp_all_reduce() lets the decoder layer suppress the module's collective when it intends to fuse it with the following norm.

vLLM reaches the same fusion from the other direction, as a compiler pass rather than a runtime flag:

class SequenceParallelismPass(VllmPatternMatcherPass):
    """
    This pass enables sequence parallelism for models.
    It identifies patterns where an AllReduce operation is followed by
    an RMSNorm (or RMSNorm and then Quantization) operation.
    These patterns are replaced with a ReduceScatter operation, followed by
    a local RMSNorm/Quantization, and then an AllGather operation.

    The general transformation is:
    Input -> AllReduce -> RMSNorm -> Output
    becomes
    Input -> ReduceScatter -> RMSNorm -> AllGather -> Output

    While this pass itself does not directly yield performance improvements,
    it lays the groundwork for subsequent fusion passes, such as
    GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can
    significantly reduce communication overhead and improve overall model
    performance.
# ...

Note what that does to the "RMSNorm is replicated" claim above: the pass exists precisely to stop it being true — sharding the norm along the token axis so each rank normalises $n/p$ rows, which costs nothing extra because the all-reduce was already going to be split into a reduce-scatter and an all-gather. It is the same total wire volume with the redundant per-rank normalisation deleted.

Two things gate it, and both matter for reading the default path correctly. It is on by default: enable_sp resolves to IS_DENSE at optimization level 2 (vllm/config/vllm.py:L303-L309), and -O2 is the default (vllm/config/vllm.py:L435), so a dense model at TP>1 gets the pass registered (vllm/compilation/passes/pass_manager.py:L161-L164). But is_applicable returns compile_range.start >= self.min_token_num (vllm/compilation/passes/fusion/sequence_parallelism.py:L604-L618), where min_token_num comes from a per-GPU-megabyte heuristic (get_sequence_parallelism_threshold, :L55-L104). A decode step is far below that threshold, so the collective this chapter counts — a plain all-reduce into a replicated norm — is what actually runs at decode; the reduce-scatter rewrite is a prefill-shaped optimisation. It is also fullgraph-only, asserted in the same method: "SequenceParallelismPass requires full-graph compilation".

Sharding at load time

The loader below narrows the incoming tensor along the output dimension. This does not prove that only shard bytes were loaded or that peak host memory equals the shard: a narrow view can retain its parent's storage, and mmap, checkpoint format and loader implementation determine RSS. Measure peak host and device allocation for the actual loader. The simple $131/\mathrm{TP}$ weight budget also omits small replicated parameters and K/V projection replication once TP exceeds the KV-head count.

    def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
        output_dim = getattr(param, "output_dim", None)

        is_sharded_weight = getattr(param, "is_sharded_weight", False)

        param_data = param.data
        if output_dim is not None and not is_sharded_weight:
            shard_size = param_data.shape[output_dim]
            start_idx = self.tp_rank * shard_size
            loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

the row loader narrows along the input dimension (vllm/model_executor/layers/linear.py:L1614-L1622, identical but with input_dim), and the QKV loader is the one place where replication becomes visible in the loader itself:

            is_sharded_weight = getattr(param, "is_sharded_weight", False)
            param_data = param_data.narrow(output_dim, shard_offset, shard_size)
            if loaded_shard_id == "q":
                shard_rank = self.tp_rank
            else:
                shard_rank = self.tp_rank // self.num_kv_head_replicas
            start_idx = shard_rank * shard_size

            if not is_sharded_weight:
                loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

At TP=8 on Llama-3-70B, num_kv_head_replicas == 1 and shard_rank == tp_rank. At TP=16 it is 2, so ranks 0 and 1 both compute shard_rank = 0 and both read the same slice of k_proj and v_proj — while reading different slices of q_proj. That single integer division is where the KV replication of Figure 2 is physically created.

Follow the number through, function by function, for rank 9 of a TP=16 Llama-3-70B:

1. LlamaAttention.__init__ (vllm/model_executor/models/llama.py:L140-L153) computes num_heads = 64 // 16 = 4 and num_kv_heads = max(1, 8 // 16) = 1. The tp_size % total_num_kv_heads == 0 branch is taken: 16 % 8 = 0, so the assert passes.
2. QKVParallelLinear.__init__ (vllm/model_executor/layers/linear.py:L1022-L1030) takes the tp_size >= total_num_kv_heads branch, sets num_kv_heads = 1 and num_kv_head_replicas = 16 // 8 = 2.
3. ColumnParallelLinear.__init__ (vllm/model_executor/layers/linear.py:L454-L473) sets output_size_per_partition. Output sizes are $[4{\cdot}128{\cdot}16,\ 1{\cdot}128{\cdot}16,\ 1{\cdot}128{\cdot}16] = [8192, 2048, 2048]$, divided by 16 to $[512, 128, 128]$ — a per-rank qkv_proj.weight of [768, 8192].
4. Loading: the WeightsMapper — the checkpoint-name-to-module rewriting table that §8.4 owns — at vllm/model_executor/models/llama.py:L345-L354 turns ...k_proj.weight into (qkv_proj, "k"), and QKVParallelLinear.weight_loader computes shard_rank = 9 // 2 = 4, start_idx = 512, and narrows rows 512:640 out of the checkpoint's [1024, 8192]. Rank 8 computes 8 // 2 = 4 and narrows the same rows.
5. At runtime, LlamaAttention.forward (vllm/model_executor/models/llama.py:L221-L231) splits the fused output with qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) = [512, 128, 128], applies RoPE, calls self.attn(q, k, v), and hands the result to self.o_proj.
6. RowParallelLinear.forward (vllm/model_executor/layers/linear.py:L1641-L1668) multiplies the rank's [512, 8192] slice of o_proj and calls tensor_model_parallel_all_reduce. That is all-reduce number one of the layer.
7. KV-cache sizing sees the same 1: SGLang's ModelConfig.get_num_kv_heads (python/sglang/srt/configs/model_config.py:L1166-L1177) returns max(1, 8 // 16) = 1, which pool_configurator._compute_cell_size (python/sglang/srt/model_executor/pool_configurator.py:L326-L333) multiplies by (head_dim + v_head_dim) * layers * kv_size = $256 \times 80 \times 2 = 40{,}960$ bytes. vLLM's ModelConfig.get_num_kv_heads (vllm/config/model.py:L1517-L1537) returns the same 1 by the same max(1, ...).

Ranks 8 and 9 now hold byte-identical K and V pages for every token in the batch, forever. That is the 1.14× in the table.

TP that does not divide the head count. The friendliest of the failures, because it is checked in config validation with a real message:

        total_num_attention_heads = self.model_arch_config.total_num_attention_heads
        tensor_parallel_size = parallel_config.tensor_parallel_size
        if total_num_attention_heads % tensor_parallel_size != 0:
            raise ValueError(
                f"Total number of attention heads ({total_num_attention_heads})"
                " must be divisible by tensor parallel size "
                f"({tensor_parallel_size})."
            )

TP that does not divide the FFN width. Less friendly. Llama-3-70B's intermediate is 28,672, which divides by 1, 2, 4, 8, 16 and 32 but not by 6 or 12. The check is a bare assert inside the layer — assert all(output_size % self.tp_size == 0 for output_size in output_sizes) (vllm/model_executor/layers/linear.py:L689) — and the underlying divide() helper raises "{} is not divisible by {}" (vllm/distributed/utils.py:L53-L57), which tells you the two numbers and nothing about which layer produced them. Search the shape.

The TP=8→16 capacity cliff. The war story of this chapter. A team moves a 70B deployment from one node to two, expecting max-concurrency to double, and gets 14%. The number to watch is what the engine logs at startup — "GPU KV cache size: %s tokens, Maximum concurrency for %s tokens per request: %.2fx" (vllm/v1/core/kv_cache_utils.py:L1925-L1931). If it barely moved when you doubled TP, you crossed $h_{kv}$. Read num_key_value_heads out of the checkpoint's config.json before you buy the second node.

Silently divergent ranks. Because every rank runs the sampler on identical logits, any per-rank divergence — a different seed, a non-deterministic kernel, an fp32 reduction order — produces ranks that pick different tokens and then compute different KV, with no error and steadily worsening output. The symptom is a model that is coherent at TP=1 and subtly wrong at TP=8. The tell is that it starts fine and degrades over the generation.

The bias double-count. Covered in §3; if a biased model is wrong at TP>1 only, check whether a custom layer replicated the bias into every rank's GEMM.

Read the sharding out of the engine rather than trusting this chapter. Both ColumnParallelLinear and RowParallelLinear implement extra_repr that prints the partitioned sizes (vllm/model_executor/layers/linear.py:L595-L601 and L1669-L1675), so printing the model prints the shard map:

python -c "
from transformers import AutoConfig
c = AutoConfig.from_pretrained('meta-llama/Meta-Llama-3-70B')
for tp in (1,2,4,8,16,32):
    q  = c.num_attention_heads // tp
    kv = max(1, c.num_key_value_heads // tp)
    rep = max(1, tp // c.num_key_value_heads)
    cell = kv * 2 * (c.hidden_size // c.num_attention_heads) * c.num_hidden_layers * 2
    print(f'TP={tp:3d}  q/rank={q:3d}  kv/rank={kv}  replicas={rep}  KV cell={cell/1024:.0f} KiB/token')
"

Then confirm it against a live engine and watch the KV-pool line move (and stop moving):

# vLLM: run separately with TP=2, then 4, then 8.
# Wait for readiness; record cache capacity and complete the benchmark.
# Ctrl-C and confirm all workers exit before starting the next configuration.
TP=2
vllm serve meta-llama/Meta-Llama-3-70B --tensor-parallel-size "$TP" \
    --max-model-len 8192

# After vLLM has fully stopped: SGLang, same knob, aliased
python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-70B \
    --tp-size 8 --context-length 8192

SGLang's flag lives at python/sglang/srt/server_args.py:L1025-L1032, declared as tp_size with the alias --tensor-parallel-size; vLLM's is tensor_parallel_size at vllm/config/parallel.py:L124-L125. The full speedup-versus-TP measurement, with a profile attributing the gap to collectives, is Lab 07 — tensor-parallel scaling.

• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019) — the paper that introduced the column→row pairing. Both engines' linear.py descends from its reference implementation, comments and all.
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models (2022) — where sequence parallelism (the reduce-scatter/all-gather rewrite of the norm) is derived. Read it alongside vllm/compilation/passes/fusion/sequence_parallelism.py.
• Pope et al., Efficiently Scaling Transformer Inference (2022) — the partitioning-strategy analysis for inference specifically, including the batch-size regimes where different layouts win. The closest thing to a first-principles treatment of Figure 3.
• NVIDIA, NVLink and NVSwitch and Quantum-2 InfiniBand — the bandwidth figures cited in §6. All interconnect numbers in this chapter are from these pages; none are measured.
• The two linear.py files read side by side — vllm/model_executor/layers/linear.py against python/sglang/srt/layers/linear.py. The clearest way to see which parts of the design are inherited from Megatron-LM and which are each engine's own; diff the two RowParallelLinear.forward bodies first.
• Neighbouring chapters: §5.2 for what to do past the node boundary, §5.3 for DP-attention (the answer to replicated MLA caches), §5.4 for the latency term $\lambda$ and custom all-reduce, and §3.5 for the GQA-versus-MLA capacity law this chapter's head arithmetic feeds.