Tensor-parallel linear layers: ColumnParallelLinear / QKVParallelLinear / RowParallelLinear

源码版本v0.25.1

Responsibilities

vLLM's tensor parallelism (TP) splits a single nn.Linear across multiple GPUs, each computing a slice, with all-gather / all-reduce on the forward and backward passes to reassemble the full result. This abstraction lives in three core classes in vllm/model_executor/layers/linear.py: ColumnParallelLinear shards along the output dim, RowParallelLinear shards along the input dim, and QKVParallelLinear handles the three Q/K/V segments in attention that have different head counts (especially under GQA / MQA). They all inherit from LinearBase, and the actual GEMM call is delegated to quant_method.apply (linear.py:551-569), so the same parallel logic runs against unquantized, fp8, AWQ, GPTQ and other backends.

When each TP linear layer is instantiated, it computes this rank's tp_rank / tp_size, divides output_size by tp_size to get output_size_per_partition (linear.py:437-442), and writes this partition info onto the parameter's output_dim / input_dim attributes; the weight_loader later uses this attribute to decide which slice to take. On top of that, QKVParallelLinear handles GQA KV head replication — when tp_size >= total_num_kv_heads, each rank holds a full copy of the KV head (num_kv_head_replicas = tp_size / total_num_kv_heads) (linear.py:994-1001).

Design motivation

Why introduce both weight_loader and weight_loader_v2 on top of LinearBase?

• Different sharding: ColumnParallelLinear shards along the output dim (each rank takes a slice of [A_1, A_2, ...]), while RowParallelLinear shards along the input dim (A = [A_1; A_2; ...], X = [X_1, X_2, ...]). So the two weight_loaders narrow along output_dim or input_dim respectively (linear.py:530-533)(linear.py:1650-1653).
• v2 goes through parameter objects: The newer weight_loader_v2 wraps the tensor in a BasevLLMParameter and calls param.load_column_parallel_weight / param.load_row_parallel_weight directly, encapsulating "which slice to take" inside the parameter object (linear.py:543-549)(linear.py:1663-1670). A quantization method can choose v2 (WEIGHT_LOADER_V2_SUPPORTED) or the legacy path.
• Q/K/V three-way fusion: QKVParallelLinear concatenates Q/K/V weights along the output dim into one big matrix; output_sizes = [q, k, v] computes divide(..., tp_size) per segment (linear.py:1003-1008). MergedColumnParallelLinear is its general form — any number of column-parallel layers can be fused.
• GQA KV replication: In QKVParallelLinear.__init__, if tp_size >= total_num_kv_heads: self.num_kv_heads = 1; self.num_kv_head_replicas = divide(tp_size, total_num_kv_heads) (linear.py:996-1001) ensures each rank can compute a complete KV head, so no extra synchronization is needed on forward.
• Shard metadata: ColumnParallelLinear's weight_loader_v2 also recognizes shard_id ("q" / "k" / "v" inside QKV) and decides which offset in the param to write to based on the shard id (linear.py:1023-1047).
• Bias only on rank 0: In RowParallelLinear.forward, bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias (linear.py:1685-1688); adding bias after reduce would double-count, so it is only added on rank 0.
• TP can be disabled: With disable_tp=True, tp_rank=0 / tp_size=1, and the layer degrades into a plain Linear — useful for non-parallel experts in MoE (linear.py:438-439).

Key files

• LinearMethodBase:141-180 — abstract base class for quantization methods, defines create_weights / apply / process_weights_after_loading.
• LinearBase:231-288 — parent class of all linear layers; update_param_tp_status writes tp_rank / tp_size onto every BasevLLMParameter.
• ColumnParallelLinear.__init__:423-491 — shards along the output dim, output_size_per_partition = divide(output_size, tp_size).
• ColumnParallelLinear weight_loader:520-549 — v1 narrows along output_dim, v2 delegates to param.load_column_parallel_weight.
• ColumnParallelLinear.forward:551-569 — quant_method.apply → optional tensor_model_parallel_all_gather.
• QKVParallelLinear.__init__:942-1021 — handles GQA KV head replication, concatenates Q/K/V output_sizes into output_size.
• QKVParallelLinear shard mapping:1023-1047 — _get_shard_offset_mapping / _get_shard_size_mapping, telling weight_loader which segment of the param to write.
• QKVParallelLinear fused checkpoint:1049-1097 — for already-fused QKV on disk (e.g. Phi-3), auto-splits then calls weight_loader_v2 per segment.
• RowParallelLinear.__init__:1572-1626 — shards along the input dim, input_size_per_partition = divide(input_size, tp_size).
• RowParallelLinear.forward:1672-1698 — if input isn't already parallel, splits first; after GEMM does tensor_model_parallel_all_reduce.
• MergedColumnParallelLinear:580-636 — fusion of multiple column-parallel layers (general form of QKV); output_sizes list decides each segment.

Data flow

Below is the core logic of ColumnParallelLinear.weight_loader — given the full tensor on disk, it narrows out the slice this rank needs at offset tp_rank * shard_size, then copy_s it into the parameter:

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)
    use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
    # bitsandbytes loads the weights of the specific portion
    # no need to narrow
    is_sharded_weight = is_sharded_weight or use_bitsandbytes_4bit

    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)

    # Special case for loading scales off disk, which often do not
    # have a shape (such as in the case of AutoFP8).
    if len(loaded_weight.shape) == 0:
        loaded_weight = loaded_weight.reshape(1)

    assert param_data.shape == loaded_weight.shape
    param_data.copy_(loaded_weight)

(linear.py:520-541) RowParallelLinear.weight_loader mirrors this, only swapping the narrow dim from output_dim to input_dim (linear.py:1650-1653). Both functions are indirectly invoked by DefaultModelLoader through model.load_weights while loading weights — every model implementation has something like loaded_weights = param.weight_loader(param, loaded_weight).

On forward, ColumnParallelLinear does all-gather (when gather_output=True, it stitches each rank's output_parallel into the full output), and RowParallelLinear does all-reduce (when reduce_results=True, it sums the sharded GEMM results). Pairing the two lets you skip one synchronization: QKVParallelLinear (column) → attention → o_proj (row) is the classic column-then-row combination.

Boundaries and failure modes

• Illegal QKV shard_id: validate_shard_id restricts the QKV shard id to "q" / "k" / "v" or None; any other value raises ValueError (linear.py:1023-1029).
• RowParallelLinear reduce_results=False + bias: Not reducing while adding bias inside the GEMM means every rank adds a copy, producing wrong results; the code raises ValueError (linear.py:1622-1626).
• FP8 block shape mismatch: _maybe_allow_fp8_block_shape_mismatch relaxes allow_fp8_block_shape_mismatch when weight_block_size doesn't evenly divide an output_partition_sizes (linear.py:493-516); otherwise the kernel rejects it.
• bitsandbytes skips narrow: With use_bitsandbytes_4bit, the checkpoint is already sharded per rank, so copy_ runs directly without narrowing (linear.py:524-527).
• Zero-dim scales: AutoFP8-style checkpoints store scales as zero-dim tensors; weight_loader uses reshape(1) as a fallback (linear.py:537-538).
• Row layer with input_is_parallel=False: When the input isn't sharded, RowParallelLinear.forward itself does split_tensor_along_last_dim to take this rank's slice (linear.py:1676-1682).

Summary

The three parallel linear classes form vLLM TP's "weight-sharding language"; how to actually compute is left to quant_method. The seam with Quantized layer loading is that ColumnParallelLinear.__init__ calls self.quant_method.create_weights(...), injecting a weight_loader (v1 or v2) into each parameter; the seam with DefaultModelLoader is that model.load_weights invokes weight_loader per tensor, slicing the full on-disk weight into this rank's segment.

See official docs: vLLM docs · README