Tensor and pipeline parallelism in practice

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

The idea#

A model can be too large for one GPU, or too slow on one GPU. Then the model runs on several GPUs. DeepInfer Explained introduces the five kinds: data, tensor, pipeline, context and expert parallelism, with their collectives (Parallelism↗). This page shows the two kinds that split one model copy, with small real numbers, and how a serving operator picks one.

Tensor parallelism: split every layer#

With tensor parallelism (TP), each GPU holds a slice of every weight matrix. All GPUs work on the same tokens at the same time. A linear layer computes Y = X W. You can split W in two ways:

• By output columns (a column-parallel linear layer). Each GPU computes some of the output features from the full input. The result on each GPU is a slice of Y.
• By input rows (a row-parallel linear layer). Each GPU gets a slice of the input features and multiplies it with its rows of W. The result on each GPU is a partial sum of the full Y. An all-reduce adds the partial sums.

The trick is to put the two splits one after the other. The MLP has two linear layers. The first layer is column-parallel, so each GPU gets a slice of the 14,336 intermediate features. The activation function works on each feature alone, so it needs no data from other GPUs. The second layer is row-parallel, and its input slice is exactly the slice that the GPU already has. So one MLP needs only one all-reduce, at the end.

Small numbers: one MLP with TP = 2#

Take a Llama-3.1-8B-style MLP: hidden size 4,096, intermediate size 14,336, BF16. T is the number of tokens in the step.

Each GPU holds 88.1 million MLP weights (176 MB) instead of 176.2 million (352 MB). It does half of the multiply work. The price is one all-reduce of T × 4096 × 2 bytes: 16 MiB for a step of 2,048 tokens.

Attention: split by heads#

Attention heads are independent of each other, so TP gives each GPU a group of heads. The QKV projection is column-parallel: each GPU computes the Q, K and V of its own heads. The output projection is row-parallel, with one all-reduce (vllm/model_executor/models/llama.py:165-181). So one decoder layer has two all-reduces: one after attention and one after the MLP.

With grouped-query attention, the model has fewer KV heads than query heads (Llama-3.1-8B: 32 query heads, 8 KV heads, head size 128). Two rules follow:

• The number of query heads must divide by TP. With TP = 3, 32 heads do not divide, and vLLM stops at start-up.
• If TP is larger than the number of KV heads, vLLM gives one KV head to each GPU and replicates it. With TP = 16, each of the 8 KV heads is on 2 GPUs. TP must then be a multiple of the KV head count.

The KV cache follows the heads. With TP = 2, each GPU stores 4 KV heads. That is 32 layers × 4 × 128 × 2 (K and V) × 2 bytes = 64 KiB per token, half of 128 KiB. With TP = 16, each GPU stores one replicated head: 16 KiB per token. The 16 GPUs then store 256 KiB per token together, twice the real size.

Pipeline parallelism: split the layers#

With pipeline parallelism (PP), each GPU, or each group of TP GPUs, holds a range of layers. With 32 layers and PP = 2, stage 0 has layers 0 to 15 and stage 1 has layers 16 to 31. Stage 0 also has the embedding. The last stage also has the final norm and the LM head.

A stage sends its hidden states to the next stage once per step. For Llama, these are two tensors: hidden_states and residual, each [T, 4096]. For T = 2,048 that is 32 MiB per step and stage boundary. TP instead sends 64 all-reduces per step (2 per layer) inside every group. So PP needs much less bandwidth, and it works over a slower network between nodes.

The cost of PP is the pipeline bubble. With one batch in flight, stage 1 waits while stage 0 works, and stage 0 waits while stage 1 works. Each GPU is busy only half of the time. vLLM keeps several batches in flight to fill the stages. PP also does not make one step faster: the stages still run one after the other.

Data and expert parallelism#

• Data parallelism (DP) runs full copies of the model. Each copy serves other requests. It adds throughput but does not make one model fit. See Data parallel serving: several engine cores behind one endpoint.
• Expert parallelism (EP) is for Mixture-of-Experts models (Mixture of Experts↗). It puts different experts on different GPUs, instead of a TP slice of every expert. See Expert parallelism and MoE serving.

A decision table for a serving operator#

The vLLM docs give the same rules for one model replica (docs/serving/parallelism_scaling.md:7-11). They also suggest PP when the GPU count does not divide the model, or when the GPUs have no NVLink (docs/serving/parallelism_scaling.md:23).

An everyday comparison#

Think of a team that must type a long book. With TP, every typist types a part of every page, and after each page they put the parts together. With PP, typist 1 types chapters 1 to 10 and typist 2 types chapters 11 to 20. Typist 2 can only start a book when typist 1 hands it over. With DP, each typist types a different book.

The comparison stops here: in TP, the "put together" step is an addition of numbers (an all-reduce), not a join of pages. And in serving, the pipeline gets a new batch every step, so several batches can be in different stages at the same time.

The picture#

Read each lane from top to bottom: one GPU and its slice of every weight. The only arrow between the lanes is the all-reduce at the bottom. Panel 2 gives the cost of this split for one MLP block.

In vLLM#

This is a basics page, so this part gives only pointers. The pages The tensor-parallel layers (linear.py and vocab_parallel_embedding.py) and Process groups: how vLLM slices its ranks explain the code in full.

1. The layers. LlamaMLP builds gate_up_proj as a MergedColumnParallelLinear and down_proj as a RowParallelLinear (vllm/model_executor/models/llama.py:95-111). RowParallelLinear.forward() calls the all-reduce when reduce_results is true and TP is larger than 1 (vllm/model_executor/layers/linear.py:1803-1806).
2. The pipeline stages. make_layers() builds only the layers of this PP rank. It fills the other places with PPMissingLayer, an identity module (vllm/model_executor/models/utils.py:886-919). A stage that is not the last returns IntermediateTensors (vllm/model_executor/models/llama.py:443-450).
3. The rank layout. vLLM numbers all ranks in the order DP × PP × PCP × TP, with TP last (vllm/distributed/parallel_state.py:2157-2163). So the ranks of one TP group are neighbours, and they normally sit on one node.

        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
            bias=bias,
            quant_config=quant_config,
            disable_tp=disable_tp,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            reduce_results=reduce_results,
            disable_tp=disable_tp,
            prefix=f"{prefix}.down_proj",
        )
        if hidden_act != "silu":
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
        self.act_fn = SiluAndMul()

    def forward(self, x):
        x, _ = self.gate_up_proj(x)
        x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
        x, _ = self.down_proj(x)
        return x

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.

    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

    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

        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)

vllm serve meta-llama/Llama-3.1-70B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2

What can go wrong#

• Heads do not divide. 32 query heads with TP = 3 or TP = 6 fail at start-up (vllm/config/model.py:1559-1564). Choose a TP that divides the head count, or use PP.
• The model has no PP support. PP needs a model class with SupportsPP. Otherwise vLLM raises NotImplementedError (vllm/config/model.py:1569-1576).
• TP across nodes. Each decoder layer has two all-reduces on the critical path. Over a network between nodes this is slow. Keep TP inside a node; put PP or DP across nodes (docs/serving/parallelism_scaling.md:9).
• TP larger than the KV heads. vLLM replicates KV heads (vllm/model_executor/layers/linear.py:1066-1068). The KV cache then stores each head several times, so the KV capacity grows less than the GPU count.
• An uneven pipeline. A stage with more layers is slower, and the other stages wait for it. Check the "unevenly partitioned" log line, or set VLLM_PP_LAYER_PARTITION (vllm/distributed/utils.py:220-228).
• DP when the model does not fit. DP copies the full model to each rank, so it does not help a model that is too large. Use TP or PP first.

Why it matters. The parallel layout decides how much memory is left for the KV cache, and how much time each step spends on communication. A wrong layout can cost more throughput than any other setting.