The tensor-parallel layers (linear.py and vocab_parallel_embedding.py)

Code version: vLLM @ 3e44c0d · 2026-10-22

The idea#

With tensor parallelism (TP) each GPU holds a slice of every weight. A column-parallel linear layer layer splits the output features; a row-parallel linear layer layer splits the input features and adds the partial sums with an all-reduce. The page Tensor and pipeline parallelism in practice shows this with one Llama MLP. This page explains the classes that do it.

A model file never calls a collective and never computes a slice. It builds layers such as QKVParallelLinear(hidden_size, head_size, total_num_heads, total_num_kv_heads) with full sizes (How a vLLM model file is written (llama.py)). Each layer class then does five jobs itself:

1. It asks the quantization config for a quantize method for its prefix (of a layer).
2. It computes its per-rank sizes and lets the quantize method make parameters of that size.
3. It gives each parameter a weight loader that cuts this rank's slice from a full checkpoint tensor ("Weight loading: from checkpoint files to GPU parameters").
4. In forward, it runs the GEMM through the quantize method.
5. It calls the collective that the split needs: an all-gather when gather_output is set, an all-reduce when reduce_results is set.

An everyday comparison. Think of a choir that sings one long chord. In a column split, each singer sings different notes of the chord; together they give the full chord. In a row split, every singer sings the same note at part of the volume. The hall adds the voices into one loud note. The comparison stops here. A hall adds the voices for free. On GPUs, the sum is an all-reduce over the links between the GPUs, and it costs time in every layer.

The picture#

Colour marks the KV head. A query head always lands on a rank that also holds the KV head it reads. So attention needs no data from other ranks; only o_proj adds the results.

In the code#

Step 1: LinearBase resolves the quantize method#

        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype
        self.quant_config = quant_config
        self.prefix = prefix
        self.allow_fp8_block_shape_mismatch = False
        self.quant_method: QuantizeMethodBase
        if quant_config is None:
            self.quant_method = UnquantizedLinearMethod()
        elif quant_method := resolve_quant_method(quant_config, self, prefix=prefix):
            self.quant_method = quant_method
        else:
            raise ValueError("All linear layers should support quant method.")
        self.return_bias = return_bias
        self.disable_tp = disable_tp
        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()
            )

1. Lines 290-296. Without a quantization config the method is UnquantizedLinearMethod. Else resolve_quant_method(quant_config, self, prefix=prefix) asks the config; it must return a method, or the constructor raises ValueError. The page Quantization methods as plugins explains the methods.
2. Lines 299-309. disable_tp=True makes the layer act as if TP were 1 (rank 0 of size 1). Else the rank and size come from the TP group, or from explicit arguments.

LinearBase is a PluggableLayer, so a plugin can replace a whole layer class by its registered name, for example column_parallel_linear (vllm/model_executor/layers/linear.py:434-436; CustomOp, out-of-tree ops and IR op providers). The bare LinearBase creates no weights; each subclass calls quant_method.create_weights() with its own sizes.

Step 2: ReplicatedLinear: no split#

A ReplicatedLinear layer makes the full weight on every rank: create_weights(self, input_size, [output_size], input_size, output_size, ...) (vllm/model_executor/layers/linear.py:379-387). Its loader copies the whole tensor after a size check (vllm/model_executor/layers/linear.py:404-412), and its forward needs no communication (vllm/model_executor/layers/linear.py:414-425). 62 model files use it (git grep at this commit). An example is the score head that as_seq_cls_model adds, with one output per class label (vllm/model_executor/models/adapters.py:365-372).

Step 3: ColumnParallelLinear: split the outputs#

        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
            ]

    def forward(
        self,
        input_,
    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
        bias = self.bias if not self.skip_bias_add else None

        # Matrix multiply.
        output_parallel = self.quant_method.apply(self, input_, bias)

        if self.gather_output and self.tp_size > 1:
            # All-gather across the partitions.
            output = tensor_model_parallel_all_gather(output_parallel)
        else:
            output = output_parallel

        if not self.return_bias:
            return output
        output_bias = self.bias if self.skip_bias_add else None
        return output, output_bias

1. Line 623. The quantize method multiplies the full input [T, in] with the local weight. The result is [T, out / tp_size].
2. Lines 625-629. Only with gather_output=True and TP > 1 does the layer call tensor_model_parallel_all_gather. The default is False (vllm/model_executor/layers/linear.py:473), because a row-parallel layer usually comes next.

divide() asserts that the size divides by TP; a wrong size stops the build with an AssertionError (vllm/distributed/utils.py:51-55, vllm/distributed/utils.py:122-126).

Step 4: MergedColumnParallelLinear: several outputs in one GEMM#

MergedColumnParallelLinear(input_size, output_sizes) stacks several column-parallel matrices, for example gate and up with output_sizes=[14336, 14336] (vllm/model_executor/models/llama.py:95-102). Each part must divide by TP, and each part is split on its own (vllm/model_executor/layers/linear.py:727-731). So a rank holds [gate slice; up slice], not a slice of [gate; up]. This is what lets SiluAndMul work locally.

        assert loaded_shard_id < len(self.output_sizes)

        shard_offset = sum(self.output_sizes[:loaded_shard_id])
        shard_size = self.output_sizes[loaded_shard_id]
        shard_offset //= self.tp_size
        shard_size //= self.tp_size

        if isinstance(param, BlockQuantScaleParameter):
            weight_block_size = getattr(self, "weight_block_size", None)
            shard_size, shard_offset = adjust_block_scale_shard(
                weight_block_size, shard_size, shard_offset
            )

        param.load_merged_column_weight(
            loaded_weight=loaded_weight,
            shard_id=loaded_shard_id,
            shard_offset=shard_offset,
            shard_size=shard_size,
        )

The shard id is an integer (0 for gate, 1 for up) or a tuple of consecutive ids. None means that the checkpoint already stores the tensor merged (vllm/model_executor/layers/linear.py:745-770).

Step 5: QKVParallelLinear: split by heads#

        # 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
        input_size = self.hidden_size
        self.output_sizes = [
            self.num_heads * self.head_size * tp_size,  # q_proj
            self.num_kv_heads * self.head_size * tp_size,  # k_proj
            self.num_kv_heads * self.v_head_size * tp_size,  # v_proj
        ]
        output_size = sum(self.output_sizes)

        super().__init__(
            input_size=input_size,
            output_size=output_size,
            bias=bias,
            gather_output=False,
            skip_bias_add=skip_bias_add,
            params_dtype=params_dtype,
            quant_config=quant_config,
            prefix=prefix,
            return_bias=return_bias,
            disable_tp=disable_tp,
        )

1. Line 1065. The query heads must divide by TP.
2. Lines 1066-1068. When TP is at least the number of KV heads, each rank gets one KV head, and num_kv_head_replicas = tp_size / total_num_kv_heads ranks share it. divide() asserts that TP is a multiple of the KV head count.
3. Lines 1069-1071. Else the KV heads are split, with no copies.
4. Lines 1073-1077. output_sizes are the sizes across all ranks, including copies: num_kv_heads × head_size × tp_size. With TP = 16 and 8 KV heads, the "total" K size is 16 × 128, twice the 8 × 128 of the checkpoint.

The parameter of a rank is [q | k | v] along dimension 0. The offsets are 0, num_heads × head_size and (num_heads + num_kv_heads) × head_size (vllm/model_executor/layers/linear.py:1101-1117). When it loads, the q part takes slice number tp_rank of the checkpoint tensor; k and v take slice number tp_rank // num_kv_head_replicas (vllm/model_executor/parameter.py:191-195). The page "Weight loading: from checkpoint files to GPU parameters" follows one tensor through this code.

The model splits the GEMM output with the same per-rank sizes: qkv.split([q_size, kv_size, kv_size], dim=-1) (vllm/model_executor/models/llama.py:229-230). QKVParallelLinear always uses gather_output=False (vllm/model_executor/layers/linear.py:1084).

Step 6: RowParallelLinear: split the inputs, add the outputs#

    def forward(
        self,
        input_,
    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
        if self.input_is_parallel:
            input_parallel = input_
        else:
            split_input = split_tensor_along_last_dim(
                input_, num_partitions=self.tp_size
            )
            input_parallel = split_input[self.tp_rank].contiguous()

        # 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

        if not self.return_bias:
            return output
        output_bias = self.bias if self.skip_bias_add else None
        return output, output_bias

1. Lines 1789-1795. With input_is_parallel=True (the default) the input is already this rank's slice, [T, in / tp_size]. Else the layer splits the last dimension and keeps its part.
2. Line 1800. Only rank 0 adds the bias; with TP ranks adding it, the all-reduce would add it TP times.
3. Lines 1803-1806. With reduce_results=True (the default) and TP > 1, the layer all-reduces the partial sums [T, out]. With False, the caller must reduce later, so a bias is then not allowed unless skip_bias_add is set (vllm/model_executor/layers/linear.py:1736-1740).

The per-rank sizes are the reverse of the column case: input_size / tp_size and the full output (vllm/model_executor/layers/linear.py:1692-1696). The bias is not split (vllm/model_executor/layers/linear.py:1742-1745).

Step 7: VocabParallelEmbedding and ParallelLMHead#

The embedding splits the vocabulary rows across the ranks. First it applies vocab padding: the vocabulary is rounded up to a multiple of padding_size, 64 by default (vllm/model_executor/layers/vocab_parallel_embedding.py:36, vllm/model_executor/layers/vocab_parallel_embedding.py:95-97, vllm/model_executor/layers/vocab_parallel_embedding.py:286-291). Each rank holds padded_vocab / tp_size rows (vllm/model_executor/layers/vocab_parallel_embedding.py:304-306).

        else:
            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,
            )
            output_parallel = self.quant_method.embedding(self, masked_input.long())

        output_dtype = output_parallel.dtype
        if is_fp8(output_dtype):
            # Each vocab token has one owner, so FP8 bytes can use int8 SUM.
            output_parallel = output_parallel.view(torch.int8)
        if not self.use_fused_embedding:
            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
        if self.reduce_results:
            output_parallel = (
                self.parallel_group.all_reduce(output_parallel)
                if self.parallel_group is not None
                else tensor_model_parallel_all_reduce(output_parallel)
            )
        return output_parallel.view(output_dtype)

1. Lines 542-550. get_masked_input_and_mask shifts each token id into the local row range and marks ids that other ranks own. On CUDA with an unquantized weight, a fused kernel does this instead (lines 527-536).
2. Line 557. Rows for ids of other ranks become zero.
3. Lines 558-563. The all-reduce adds the ranks: exactly one rank gives a non-zero row for each token, so the sum is the embedding.

ParallelLMHead has the same weight layout, but its forward raises: the weight is used by LogitsProcessor (vllm/model_executor/layers/vocab_parallel_embedding.py:640-642). LogitsProcessor multiplies the hidden states with the local rows, gathers the vocabulary slices across the TP group, and cuts off the padding (vllm/model_executor/layers/logits_processor.py:180-199). It uses an all-gather: Platform.use_all_gather() returns True, and no built-in platform overrides it (vllm/platforms/interface.py:1275-1278). tie_weights makes the head use the embedding's weight object (vllm/model_executor/layers/vocab_parallel_embedding.py:88-92).

Step 8: the collectives and the communicator#

The layers call two small functions. Both use the TP group of this process:

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_)


def tensor_model_parallel_all_gather(
    input_: torch.Tensor, dim: int = -1
) -> torch.Tensor:
    """All-gather the input tensor across model parallel group."""
    return get_tp_group().all_gather(input_, dim)

GroupCoordinator.all_reduce returns the input unchanged when the group has one rank. On CUDA, ROCm, XPU and TPU it calls the registered op torch.ops.vllm.all_reduce, so torch.compile sees one opaque op. Other platforms call the device communicator directly (vllm/distributed/parallel_state.py:740-746, vllm/distributed/parallel_state.py:571-573, for example vllm/platforms/cuda.py:708-709). The CUDA communicator then picks one of several all-reduce implementations, such as custom all-reduce, symmetric-memory NCCL or plain NCCL (vllm/distributed/device_communicators/cuda_communicator.py:362-446). The page Process groups: how vLLM slices its ranks explains the groups and the communicators.

Deep dive#

The sections below give the full detail. Open them one at a time, or use Open all deep dives at the top of the page.

What can go wrong#

• A TP size that does not fit the model. Start-up checks the query heads first (vllm/config/model.py:1559-1564). Other sizes fail later, while the layers are built, in divide() or an assertion. Examples are the MLP intermediate size, or a KV head count that does not fit TP (vllm/model_executor/layers/linear.py:731, vllm/model_executor/layers/linear.py:1066-1071).
• KV replication eats KV cache. With TP larger than the KV head count, each rank stores one full KV head, so the cluster stores copies (vllm/model_executor/layers/linear.py:1066-1068). TP = 16 for an 8-KV-head model gives twice the KV bytes of TP = 8 in total.
• A custom layer without a reduce. A model that sets reduce_results=False must call the all-reduce itself. Without it, each rank keeps a partial sum, and the output is wrong with no error (derived) (vllm/model_executor/layers/linear.py:1803-1806).
• TP across nodes. Every decoder layer has two all-reduces on the critical path (Tensor and pipeline parallelism in practice). Over a slow network, prefer PP or DP across nodes.

Why it matters. These layer classes decide the memory per GPU, the number of collectives per step and the rules for valid TP sizes. When you know them, you can predict the per-GPU shapes of any model and read the errors of a wrong TP size.