vLLM 源码完全拆解（十一）：模型并行与 Worker 层——Tensor Parallel 如何切分权重并在多卡间同步

一个矩阵乘法如何被切开#

张量并行的思想很简单，但要讲清楚它为什么省显存、为什么需要通信，得从一次最简单的线性层说起。

一个线性层（dense 层）做的计算是：

其中 $X$ 是输入（形状 $[b, k]$，$b$ 是 batch 里的 token 数，$k$ 是输入维度），$W$ 是权重（形状 $[k, n]$，$n$ 是输出维度），$b$ 是偏置（形状 $[n]$），$Y$ 是输出（形状 $[b, n]$）。

现在有 $t$ 张卡（vLLM 里叫 tensor_parallel_size）。张量并行要对上面这次 GEMM 做拆分，有两种基本切法，分别对应两种”并行线性层”，这正是 vLLM 源码里 ColumnParallelLinear 和 RowParallelLinear 两个类的来历。

Column（列）切法——按输出维度 $n$ 切权重。把 $W$ 沿最后一维（列方向）切成 $t$ 块，每卡拿 $W_i \in \mathbb{R}^{[k, n/t]}$。因为：

所以每张卡用相同的完整输入 $X$ 和自己的那列权重 $W_i$ 算出一块输出 $Y_i = XW_i$。所有卡算完，各占 $Y$ 的 $n/t$ 列。也就是说，权重被水平切开，输入是全量复制给每张卡，输出按列分摊。

Row（行）切法——按输入维度 $k$ 切权重。把 $W$ 沿第一维（行方向）切成 $t$ 块，每卡拿 $W_i \in \mathbb{R}^{[k/t, n]}$。把输入 $X$ 也沿列（最后一维）切成 $t$ 份 $X_i \in \mathbb{R}^{[b, k/t]}$，有：

每张卡用 $X_i$ 和 $W_i$ 算出一个部分和，最后要把 $t$ 张卡的部分和加起来才是完整输出。这意味着 Row 切法天然需要一次求和通信（reduce）。

这两种切法，一个”输出分摊、无需直接求和（因为每块列结果独立）“，一个”需要累加部分和”。它们一个像”分头写文章的各章节”，一个像”各自算一遍再合并结果”。

为什么需要 all-reduce 而不是简单地各自算完#

回到上面的分析：Column 切法看起来每张卡算完自己的列就完事了，不需要通信。但在 Transformer 里事情没这么简单——因为下一层往往需要完整的输出。

考虑一个具体的 Transformer 结构：self-attention 的 qkv_proj 之后，下一层是注意力计算，然后接 o_proj。如果你用 Column 切 qkv_proj，那么每个 head 的 Q/K/V 分片在各自的卡上齐了（因为注意力是按 head 分的，head 恰好就被 Column 切分到各卡），可以直接在本卡做注意力。但注意力算完的 attn_output 是一个完整向量，o_proj 的输入需要完整 hidden state。

于是 vLLM 的典型做法是”列行交替”：qkv_proj 用 Column（列）切，o_proj 用 Row（行）切。o_proj 的 Row 切需要把输入（即上一层的注意力输出）按列切给各卡、然后再 all-reduce。这样在 qkv/o 之间插一次 all-reduce，attn 的输出被拼回完整的 hidden state 喂给下一层。

同样，FFN 里 gate_up_proj 用 Column（列）切、down_proj 用 Row（行）切，down_proj 之后做一次 all-reduce。这就是经典的”MLP 中间 all-reduce + attention 输出 all-reduce”,也就是 Megatron-LM 那套”f-g 交替（f forward / g backward）“模式的推理侧版本。

所以结论是：张量并行的”切”本身不产生通信，但为了把结果拼回完整维度喂给下一层，每一对 Column–Row 之间必须有一次 all-reduce。通信的代价就藏在这一次次 all-reduce 里。

MLP 的张量并行切分方式（来源：Megatron-LM 论文 Fig.3）。左侧 Column 切 gate_up（f 阶段），右侧 Row 切 down（g 阶段），中间用 all-reduce（论文里标为 f / g 的通信）把部分和叠加成完整输出。vLLM 的 MergedColumnParallelLinear + RowParallelLinear 组合正是这套结构的推理侧实现。

Self-Attention 的张量并行切分（来源：Megatron-LM 论文 Fig.5）。QKV 投影沿 head 维度切（Column），注意力输出投影（o_proj）做 all-reduce（Row）。vLLM 的 QkvParallelLinear（Column）+ o_proj（Row）严格对应此结构。

现在你理解了张量并行的宏观思路。接下来深入到 vLLM 源码，看这些线性层到底怎么在构造时算”每卡拿哪一块”，又怎么在 load_weights 阶段只把自己那块切片填进来。

GroupCoordinator：通信的本体#

在讲线性层之前，必须先认识 vLLM 所有通信的载体——GroupCoordinator（vllm/distributed/parallel_state.py）。它是一个对 PyTorch ProcessGroup 的封装，统一管一组进程之间的通信。

它的构造函数接收 group_ranks: list[list[int]]，也就是”把哪些全局 rank 划成一组”。它对每组新建两个通信组（第 223-234 行）：

1
for ranks in group_ranks:
2
    device_group = torch.distributed.new_group(
3
        ranks, backend=torch_distributed_backend)   # 设备通信组，GPU 上用 NCCL
4
    cpu_group = torch.distributed.new_group(ranks, backend="gloo")  # CPU 通信组

这里的关键设计：同一个进程组，同时维护两个后端通信。device_group 用 NCCL（GPU 之间高速通信，走 NVLink/IB），cpu_group 用 Gloo（CPU 之间同步、广播元数据）。为什么要有 CPU 组？因为有些信息（比如 broadcast 一个 dict、同步一个标志位）不需要 GPU 显存里的张量参与，用 CPU 组更轻量、也更不容易和计算流冲突。

GroupCoordinator 对外暴露的核心 API 就是那几个通信原语，比如 all_reduce（第 336 行）：

1
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
2
    if self.world_size == 1:
3
        return input_               # 单卡直接返回，省一次通信
4
    if self.use_custom_op_call:
5
        return torch.ops.vllm.all_reduce(input_,
6
                                         group_name=self.unique_name)
7
    else:
8
        return self._all_reduce_out_place(input_)

以及 all_gather、reduce_scatter、broadcast、send、recv 等。注意这里有个细节：vLLM 用 torch.ops.vllm.all_reduce 这个自定义 op，传入的是 group_name（字符串）而不是 group 对象。原因是这套代码要能被 torch.compile/Dynamo 用，而 Dynamo 不支持把任意对象（self）传给自定义 op，所以用一个字符串名字做间接寻址。

Tensor 并行组（TP group）怎么生成#

光有 GroupCoordinator 还不够，还得知道哪些 rank 组成”张量并行组”。这是 initialize_model_parallel 干的事（parallel_state.py 第 977 行起）。它的核心是把 world 里所有 rank 排成一个 4 维矩阵，然后按不同维度切片得到不同的组：

1
# 布局顺序: ExternalDP x DP x PP x TP
2
all_ranks = torch.arange(world_size).reshape(
3
    -1, data_parallel_size, pipeline_model_parallel_size,
4
    tensor_model_parallel_size)
5

6
# 张量并行组: 按最后一维(TP)切
7
group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
8
_TP = init_model_parallel_group(group_ranks, ...)
9

10
# 流水线并行组: 转置后按 PP 维切
11
group_ranks = all_ranks.transpose(2, 3).reshape(
12
    -1, pipeline_model_parallel_size).unbind(0)
13
_PP = init_model_parallel_group(group_ranks, ...)

文档里给了一个 8 卡、TP=2、PP=4 的例子（第 991-998 行）：会生成 4 个 TP 组 [g0,g1], [g2,g3], [g4,g5], [g6,g7] 和 2 个 PP 组 [g0,g2,g4,g6], [g1,g3,g5,g7]。

一个很关键的约定是：TP 组的成员在 rank 上是连续的。[0,1]、[2,3]… 之所以这样分，是为了让同组落在同一台机器（同一个 NVLink 域）上，通信走得快。这是初始化时给用户的提示：相邻 rank 应尽量放在同一台 DGX 上。

initialize_model_parallel 结束后，全局就有几个单例组：_TP、_PP、_DP、_EP，分别通过 get_tp_group()、get_pp_group() 等拿到。模型的 o_proj、down_proj 里的 all-reduce，用的就是 get_tp_group() 这个组。

张量并行与数据并行混合的划分示意（来源：Megatron-LM 论文 Fig.8，示意了 TP 组沿模型层内部切、DP 组沿 batch 维切）。vLLM 的 initialize_model_parallel 用同样的方式，把 rank 排成矩阵后按不同维切片，分别得到 TP 组和 PP 组。

ColumnParallelLinear：沿输出维切权重#

现在进入正题。ColumnParallelLinear（vllm/model_executor/layers/linear.py 第 370 行）是张量并行的第一个主角。它在构造时就算好”本卡拿多少列”：

1
# 沿最后一维(输出维)切分权重
2
self.tp_size = get_tensor_model_parallel_world_size()
3
self.input_size_per_partition = input_size
4
self.output_size_per_partition = divide(output_size, self.tp_size)

divide 就是 output_size // tp_size，并且断言能整除。也就是说，整张权重矩阵的宽度 $n$ 被等分到 $t$ 张卡，本卡只持有 $[k, n/t]$ 这一块。这直接解释了显存怎么省：每张卡只存 1/t 的权重。

权重加载时（weight_loader，第 455 行），它要知道”从完整权重里切哪一段给我”。切的位置由 tp_rank（本卡在 TP 组里的序号）决定：

1
def weight_loader(self, param, loaded_weight):
2
    tp_rank = get_tensor_model_parallel_rank()
3
    output_dim = getattr(param, "output_dim", None)
4
    ...
5
    if output_dim is not None and not is_sharded_weight:
6
        shard_size = param_data.shape[output_dim]      # 本卡拿多少
7
        start_idx = tp_rank * shard_size               # 从完整权重的哪一格开始
8
        loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
9
    ...
10
    param_data.copy_(loaded_weight)                    # 只把这一块填进本卡参数

param.output_dim 是一个元数据，指示这个参数沿哪个维度切分（对权重矩阵是最后一维 -1）。narrow 就是 Python 端的张量切片。于是每次加载权重时，vLLM 从 huggingface 的完整 checkpoint 里 narrow 出 [start_idx : start_idx+shard_size] 这一段，填进本卡的参数。这就是”切分权重”最直接的一步代码。

ColumnParallelLinear.forward（第 503 行）：

1
def forward(self, input_):
2
    output_parallel = self.quant_method.apply(self, input_, bias)
3
    if self.gather_output:
4
        output = tensor_model_parallel_all_gather(output_parallel)
5
    else:
6
        output = output_parallel
7
    return output

这里有两个分支。gather_output=False（默认）时，本卡的输出就是自己的那 $n/t$ 列，不立即通信；这个输出通常被交给下一个 Row 层直接消费。gather_output=True 时，才用 all-gather 把各卡的列拼成完整的 $n$ 维。在标准 Transformer 里，qkv_proj 用 gather_output=False，因为注意力的 head 本来就是按卡切好的，不需要拼回全部 head。

注意 apply 是 self.quant_method 提供的——它统一了”是否量化、用哪个 kernel”这些细节。这让我们看到 vLLM 的一个干净设计：切分逻辑（线性层的 shape 计算、narrow、通信时机）和计算逻辑（quant_method）解耦。

MergedColumnParallelLinear / QkvParallelLinear：几个层拼成一个#

Transformer 里有两类”一个模块拆成几个投影”的层，vLLM 用两个子类优化它们。

MergedColumnParallelLinear（第 530 行）：适用于 FFN 的 gate_up_proj——把 gate 和 up 两个线性层沿输出维拼成一个长权重，构造时传 output_sizes=[intermediate, intermediate]。这样一次 GEMM 同时算 gate 和 up 的输出，省一次 kernel launch。

QkvParallelLinear（第 663 行左右）：作用于 self-attention 的 q/k/v 投影。它和 MergedColumn 不同，因为它要处理一个非常特殊的切分问题——多头注意力（MHA）里的 head 在 TP 下怎么分。

看它的构造逻辑（关键在 num_heads 和 num_kv_heads 的计算）：

1
tp_size = get_tensor_model_parallel_world_size()
2
self.num_heads = divide(self.total_num_heads, tp_size)     # Q 头按卡平分
3
if tp_size >= self.total_num_kv_heads:
4
    self.num_kv_heads = 1
5
    self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
6
else:
7
    self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
8
    self.num_kv_head_replicas = 1

这里体现的是 GQA（Grouped Query Attention）在 TP 下的切分。标准 MHA 里 total_num_kv_heads == total_num_heads，Q/K/V 都按卡平分。但如果是 GQA（比如 KV 头只有 8 个），且 tp_size >= num_kv_heads，那么 K/V 头不够分，于是每个 KV 头要在多张卡上复制（num_kv_head_replicas > 1）。这带来两个后果：

1. 显存不对称：Q 头被切开（每卡 1/t），但 KV 头被复制（每卡都存一份）。K/V 投影的显存没有被切省，这正对着 GQA 的”KV 缓存小”特性——KV 本来就不大，复制的代价可以接受。
2. 前向需要把 Q 头拼回完整来对齐：由于 KV 头复制在每张卡上而 Q 头只占 1/t，注意力计算时本卡的 Q 只能看到本卡的 KV，但某些 query 要 attend 的 key 在别的卡上…… 这个我们在讲 Attention 的篇章里处理，这里只需记住：GQA + TP 的关键是 KV 头按 “复制还是切分” 二选一划分，这就是 num_kv_head_replicas 的分支。

本站已有 GQA 完全拆解 专门讲 GQA 的动机，可以配合阅读。

RowParallelLinear：沿输入维切权重 + all-reduce#

RowParallelLinear（第 626 行左右）是张量并行的第二个主角，负责”切输入 + 累加部分和”。

构造时沿第一维切（第 656 行附近）：

1
self.tp_rank = get_tensor_model_parallel_rank()
2
self.input_size_per_partition = divide(input_size, self.tp_size)  # 输入维切
3
self.output_size_per_partition = output_size                       # 输出不切

forward（input_is_parallel 参数很关键）：

1
def forward(self, input_):
2
    if self.input_is_parallel:
3
        input_parallel = input_                    # 输入已经按卡切好，直接用
4
    else:
5
        # 输入还是完整的，在 forward 里切
6
        splitted_input = split_tensor_along_last_dim(input_, num_partitions=self.tp_size)
7
        input_parallel = splitted_input[self.tp_rank].contiguous()
8

9
    output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
10
    if self.reduce_results and self.tp_size > 1:
11
        output = tensor_model_parallel_all_reduce(output_parallel)   # All-reduce！
12
    else:
13
        output = output_parallel
14
    return output

input_is_parallel=True 表示”上一层的输出已经按列切好给我的卡了，不用再切”。在标准的 qkv→o 链路里，o_proj 的输入是 attention 的输出，它按 head 切在每张卡上，所以 input_is_parallel=True；而最外层的某些层（比如第一个线性层的输入是完整 embedding），需要 input_is_parallel=False 并在 forward 里先 split_tensor_along_last_dim。

最核心的一行是 tensor_model_parallel_all_reduce(output_parallel)。它调用：

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

也就是在 TP 组内做一次 all-reduce，把所有卡算的部分和累加，得到完整输出。这一下就闭环了我们前面”为什么需要 all-reduce”的解释：Row 层的每张卡算的是部分和，必须 all-reduce 才能喂给下一层。

一个典型 block 的”列行交替”拼装#

把以上拼起来，看 Llama 模型里一个 Transformer Block 的 MLP 部分（llama.py）：

1
self.gate_up_proj = MergedColumnParallelLinear(   # Column: 沿输出维切
2
    hidden_size, 2 * intermediate_size, ...)
3
self.down_proj = RowParallelLinear(               # Row: 沿输入维切 + all-reduce
4
    intermediate_size, hidden_size, ...)

自注意力部分：

1
self.qkv_proj = QkvParallelLinear(                # Column: 沿 head 维切
2
    hidden_size, head_size, num_heads, num_kv_heads, ...)
3
self.o_proj = RowParallelLinear(                  # Row: all-reduce 拼回完整 hidden
4
    hidden_size, hidden_size, ...)

这就是一个完整的张量并行单元：Column 进 → Row 出，Row 结束做 all-reduce。每层 block 都这样交替，模型的显存被均匀切到 $t$ 张卡，代价是每层多一次 all-reduce。这也解释了为什么 TP 在单机、NVLink 高带宽环境下划算——通信快，摊薄了切分带来的好处；而跨机 TP 会因 IB 带宽受限。这也是为何大模型跨机部署更倾向用「流水线并行」或「专家并行」，我们在 Worker 层和 MoE 篇里分别涉及。

通信的两种底层实现#

上面讲到 get_tp_group().all_reduce() 是 GroupCoordinator.all_reduce。它内部走 torch.ops.vllm.all_reduce(...)。为什么不用 torch.distributed.all_reduce 直接做？vLLM 自己实现了一个 Custom AllReduce（vllm/distributed/device_communicators/custom_all_reduce.py），默认开启（set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)，见 gpu_worker.py）。

Custom AllReduce 的思路是：把 NCCL all-reduce 的下半场（broadcast-reduce）用共享内存 + 一个额外 buffer 优化掉。简单说，它在 TP 组内选一个”主卡”，其它卡先把各自的分片结果通过 PCIe/NVLink 写到一个预先分配的 buffer，主卡聚合后再广播回去。相比每个 step 都做完整 NCCL all-reduce，它省掉了握手同步的固定开销，在小规模（单机、TP 卡数不多）、小数据量下收益明显。这也是 vLLM 官方 benchmark 里 TP 场景能跑出低延迟的原因之一。

不过 Custom All-Reduce 有前提：它假设同组卡在同一台机器、共享内存（多节点下无法用共享内存做聚合）。所以跨机场景会回退到标准 NCCL。这也是 disable_custom_all_reduce 这个开关存在的原因——某些环境（比如和某些通信库冲突）需要关掉。

all-gather：把分片拼回完整#

除了 all-reduce，张量并行还常用 all_gather——在 gather_output=True 的 Column 层或某些场景下，把各卡持有的分片按维度拼回完整张量。GroupCoordinator.all_gather：

1
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
2
    world_size = self.world_size
3
    if world_size == 1:
4
        return input_
5
    return torch.ops.vllm.all_gather(input_, dim, world_size,
6
                                     group_name=self.unique_name)

all-gather 的语义是：组内每张卡各持有 $[b, n/t]$ 的一部分，交换后每张卡都拿到完整的 $[b, n]$。它和 all-reduce 的区别是：all-reduce 是”各卡的部分和 → 每卡都拿到总和”，all-gather 是”各卡的分片 → 每卡都拿到完整拼接”，后者不做数值累加。

至此，“模型怎么切、通信怎么走”已经清楚。但还有一个问题没解决：这些 rank 进程是怎么被创建、怎么被指派角色、如何同时收到同一个执行指令的？这就进入第三部分——Worker 层。

一次前向的”总-分-总”三阶段#

在 vLLM V1 里，Executor.execute_model 是入口。回顾 第一篇整体架构 讲过的分层：EngineCore 持有 Executor，Executor 持有若干 Worker，每个 Worker 进程持有一个 GPUModelRunner。一次 GPU 前向是三阶段：

1. 总（广播）：Executor 把 SchedulerOutput（本次要执行的 batch 的调度结果）通过广播消息队列发给所有 Worker。
2. 分（并行执行）：每个 Worker 的 GPUModelRunner 在自己的 GPU 上执行模型前向，期间通过 TP 组 all-reduce 等在卡间同步。TP=1 时各卡独立；TP>1 时同 PP stage 内的卡协作。
3. 总（回收）：只有指定的”输出 rank”（通常是 TP 组 rank 0 和最后一个 PP stage 的 rank 0）把 ModelRunnerOutput 通过响应消息队列发回 Executor，其余 Worker 返回 None。

这个”广播-并行-回收”结构是理解 Worker 层的钥匙。

MultiprocExecutor：单机多进程的执行器#

vLLM V1 默认的多卡执行器是 MultiprocExecutor（vllm/v1/executor/multiproc_executor.py）。它用一个消息队列广播 SchedulerOutput，为每个 rank 创建一个 Worker 子进程，然后每个进程跑 worker_busy_loop 循环。注意 Executor.get_class（abstract.py）根据 distributed_executor_backend 选择实现："mp" → MultiprocExecutor、"ray" → RayDistributedExecutor、"uni" → 单进程 UniProcExecutor。

_init_executor（第 43 行）的建进程部分：

1
for rank in range(self.world_size):
2
    unready_workers.append(
3
        WorkerProc.make_worker_process(
4
            vllm_config=self.vllm_config,
5
            local_rank=rank, rank=rank,
6
            distributed_init_method=distributed_init_method,
7
            input_shm_handle=scheduler_output_handle,
8
        ))
9
self.workers = WorkerProc.wait_for_ready(unready_workers)

world_size = tensor_parallel_size * pipeline_parallel_size。每个 rank 一个进程，local_rank=rank（单机场景 local_rank 等于全局 rank），这就把每张卡绑定到一个 Worker 进程。

collective_rpc（第 166 行）是广播的核心：

1
def collective_rpc(self, method, timeout=None, args=(), kwargs=None,
2
                   non_block=False, unique_reply_rank=None):
3
    self.rpc_broadcast_mq.enqueue((send_method, args, kwargs, unique_reply_rank))
4
    workers = (self.workers[unique_reply_rank],) if unique_reply_rank is not None \
5
              else self.workers
6
    responses = []
7
    for w in workers:
8
        result = get_response(w, dequeue_timeout) if not non_block \
9
                 else self.io_thread_pool.submit(get_response, w, dequeue_timeout, self.shutdown_event)
10
        responses.append(result)
11
    return responses

要点：

• 广播：rpc_broadcast_mq.enqueue(...) 把一个 (method, args, kwargs, unique_reply_rank) 元组广播给所有 Worker（共享内存环形消息队列，见 shm_broadcast.py）。
• 选择性回传：unique_reply_rank 指定”只要哪个 rank 回结果”。比如 execute_model 只想要最后一个 PP stage 的 TP rank 0 回 ModelRunnerOutput，于是 unique_reply_rank = self._get_output_rank()：

  1
  def _get_output_rank(self):
  2
      return self.world_size - self.parallel_config.tensor_parallel_size

  这个例子是文档注释里给的：TP=8、PP=4 → world_size=32，world_size - tp_size = 32 - 8 = 24，正好是 PP rank=3（最后一 stage）且 TP rank=0 的那张卡。

execute_model（第 154 行）：

1
def execute_model(self, scheduler_output):
2
    (output, ) = self.collective_rpc(
3
        "execute_model", args=(scheduler_output, ),
4
        unique_reply_rank=self.output_rank,
5
        non_block=self.max_concurrent_batches > 1,
6
        timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS)
7
    return output

non_block 是流水线并行（PP>1）时需要的：PP 有多个 stage，可以异步把 batch 交出去，用 io_thread_pool（单线程）去异步取结果。这也对应 max_concurrent_batches = pipeline_parallel_size（第 273 行）——PP 下允许多个 batch 同时在流水线里。

WorkerProc 与 worker_busy_loop：Worker 进程的主循环#

每个 Worker 进程的入口是 WorkerProc.worker_main → 构造 WorkerProc → 调用 worker.worker_busy_loop()。WorkerProc.__init__ 里做的是”初始化这个进程对应的 worker 并加载模型”：

1
wrapper = WorkerWrapperBase(vllm_config=vllm_config, rpc_rank=rank)
2
wrapper.init_worker(all_kwargs)          # 创建 Worker（按 rank 区分 driver/non-driver）
3
...
4
self.worker.init_device()                # 初始化 CUDA 设备 + 分布式环境
5
self.worker.load_model()                 # 构造 GPUModelRunner 并加载权重

然后主循环 worker_busy_loop（第 507 行）：

1
def worker_busy_loop(self):
2
    while True:
3
        method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue()
4
        try:
5
            if isinstance(method, str):
6
                func = getattr(self.worker, method)     # 按名字调用 worker 的方法
7
            elif isinstance(method, bytes):
8
                func = partial(cloudpickle.loads(method), self.worker)  # 任意 callable
9
            output = func(*args, **kwargs)
10
        except Exception as e:
11
            ...
12
            if output_rank is None or self.rank == output_rank:
13
                self.worker_response_mq.enqueue((ResponseStatus.FAILURE, str(e)))
14
            continue
15
        if output_rank is None or self.rank == output_rank:
16
            self.worker_response_mq.enqueue((ResponseStatus.SUCCESS, output))

这就是”广播的执行体”：每个 Worker 从广播队列取到同一个 ("execute_model", (scheduler_output,), {}, rank24)，调用 self.worker.execute_model(scheduler_output)，只有 self.rank == output_rank（即 rank 24）的进程把结果写回响应队列，其余进程算完就丢弃。每个进程执行的是同一个方法、同一批调度数据，但各自只在自己的 GPU 上跑自己的模型分片——这就是张量并行”同时执行、各算各的、通信同步”的落点。

另外 WorkerProc 有一个 is_driver_worker 的概念（第 333 行）：

1
is_driver_worker = (rank % vllm_config.parallel_config.tensor_parallel_size == 0)

每个 TP 组的第一个 rank（rank % tp_size == 0）是 driver worker。driver worker 负责一些”只需一份”的工作（比如采样结果的收集）。这与我们前面”最终只回传输出 rank”的设计呼应——同 TP 组内，通常只让 driver 回传结果。

Worker / GPUWorker：单卡上干了什么#

每个 Worker 进程里跑的是 GPUWorker（vllm/v1/worker/gpu_worker.py）。它的 init_device（第 120 行）在每张卡上做三件事：

1. 绑定设备：self.device = torch.device(f"cuda:{self.local_rank}")，torch.cuda.set_device(self.device)——把当前进程绑定到 local_rank 这张卡。
2. 初始化分布式环境：init_worker_distributed_environment(...)，内部调用 init_distributed_environment(world_size, rank, distributed_init_method, local_rank, backend)，然后 ensure_model_parallel_initialized(tp_size, pp_size) 确保 TP/PP 组建好。
3. 构造 GPUModelRunner：self.model_runner = GPUModelRunner(self.vllm_config, self.device)。

execute_model（第 298 行）展示了 PP 下 stage 之间的 send/recv：

1
def execute_model(self, scheduler_output):
2
    intermediate_tensors = None
3
    if not get_pp_group().is_first_rank:                 # 非第一个 PP stage
4
        intermediate_tensors = IntermediateTensors(
5
            get_pp_group().recv_tensor_dict(all_gather_group=get_tp_group()))
6
    output = self.model_runner.execute_model(scheduler_output, intermediate_tensors)
7
    if parallel_config.distributed_executor_backend != "external_launcher" \
8
       and not get_pp_group().is_last_rank:              # 非最后一个 PP stage
9
        get_pp_group().send_tensor_dict(output.tensors, all_gather_group=get_tp_group())
10
        return None
11
    return output if self.is_driver_worker else None

这里体现了两件事：

• PP 的 stage 间传递：get_pp_group().recv_tensor_dict(...) 从上一个 PP stage 接收中间张量，算完后 send_tensor_dict(...) 发给下一个 stage，最后一个 stage 才产生 ModelRunnerOutput。这和 TP 的 all-reduce 不同，属于流水线并行（PP）特有的点对点传输。本文主角是 TP，PP 在这里只需理解”每个 PP stage 是一整层模型的子序列，stage 之间用 send/recv 串成流水线”。
• TP 内 only driver 回传：return output if self.is_driver_worker else None——同 TP 组里只有 driver worker 返回结果，其它返回 None。这样 Executor 拿到的输出是唯一的。

注意 recv/send_tensor_dict 传了 all_gather_group=get_tp_group()，表示在 PP 的 send/recv 前可能需要先在 TP 组内 all-gather，因为 stage 的输入可能是 TP 分片。这些细节在 vLLM 的 PP 实现里有一套完整的 schedule 机制，属于更深的范畴，本篇点到为止。

为什么 Worker 层和模型层是解耦的#

看到这里你会发现 vLLM 的结构非常”正交”：

• 模型层（linear.py）负责”怎么切、怎么同步”，它只关心 TP 组的通信原语，不关心进程怎么起。
• Worker 层（worker/）负责”每张卡一个进程、进程里跑什么”，它只关心分发执行，不关心模型内部怎么切。

两者的接口是 GPUModelRunner.execute_model。Executor 说”这个 batch 你跑一下”，Worker 说”我把权重切好、把 TP 组建好、在自己的卡上跑，然后 all-reduce”。这种解耦让 vLLM 能轻松切换执行后端（mp / ray / uni）、切换设备（GPU/TPU/XPU，各有自己的 model_runner），而模型层的切分逻辑完全复用。