vLLM Code Reading (Section 5): GPUWorker.execute_model 与 PP/TP 通信

GPUWorker.execute_model() 是 EngineCore 广播 execute_model RPC 之后，单个 worker rank 上进入 GPU 执行链路的顶层入口。它承接上一节的 EngineCore -> Worker 控制面，把一次 SchedulerOutput 转交给 GPUModelRunner，同时处理 pipeline parallelism 的边界通信。

1. 代码入口

GPUWorker.execute_model() 本身不直接跑 Transformer layer，而是负责把一次 SchedulerOutput 转交给 GPUModelRunner，同时处理 pipeline parallelism 的边界通信。

这段函数可以先按职责拆成一句话：GPUWorker 管 PP 边界，GPUModelRunner 管本 rank 的模型执行。 也就是说，GPUWorker.execute_model() 解决的是“本 PP stage 要不要收上游 hidden states、要不要把本 stage 结果发给下游”，而不是“attention 或 MLP 具体怎么切权重”。

核心路径如下：

flowchart TD
    A["GPUWorker.execute_model(scheduler_output)"] --> B["wait previous _pp_send_work"]
    B --> C["prepare local execution flags"]
    C --> D{"PP + SP + forward?"}
    D -- yes --> E["compute all_gather_tensors"]
    D -- no --> F["skip SP-specific prep"]
    E --> G{"not PP first rank?"}
    F --> G
    G -- yes --> H["irecv_tensor_dict from previous PP stage"]
    G -- no --> I["intermediate_tensors = None"]
    H --> J["wrap AsyncIntermediateTensors"]
    I --> K["model_runner.execute_model"]
    J --> K
    K --> L{"output type?"}
    L -- "ModelRunnerOutput / Async / None" --> M["return output"]
    L -- "IntermediateTensors" --> N["isend_tensor_dict to next PP stage"]
    N --> O["save handles in _pp_send_work"]
    O --> P["return None"]

这张图里最重要的分界点是 model_runner.execute_model(...)。在它之前，GPUWorker 主要处理通信准备；在它之后，GPUWorker 根据返回值决定是直接返回给 executor，还是把 intermediate tensors 继续发给下一个 PP stage。

2. 确保上一轮 PP 发送完成

函数一开始先处理上一轮留下的非阻塞 PP send：

if self._pp_send_work:
    for handle in self._pp_send_work:
        handle.wait()
    self._pp_send_work = []

这是 PP 通信里非常关键的生命周期管理。非最后一个 PP stage 在函数末尾会通过 isend_tensor_dict(...) 把 intermediate tensors 发给下一个 stage；这个 send 是 non-blocking 的，因此函数返回时传输可能还没结束。

下一轮执行开始前必须等待这些 send 完成，否则可能出现两个问题。第一，上一轮通信还在读 buffer，当前轮已经复用或覆盖了这些 tensor；第二，pipeline stage 之间的消息顺序被破坏。_pp_send_work 保存的就是上一轮异步发送句柄，本轮开始时统一 wait() 并清空。

因此，这一块不是性能装饰，而是保证 PP 通信正确性的同步点。

3. 准备本轮执行状态

接下来函数初始化本轮局部状态：

intermediate_tensors = None
forward_pass = scheduler_output.total_num_scheduled_tokens > 0
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
all_gather_tensors = {}
compilation_config = self.vllm_config.compilation_config
parallel_config = self.vllm_config.parallel_config

这里最重要的是 forward_pass。不是每次 execute_model 都一定有真实 forward。例如某些 step 可能只是释放 request、处理 finished 状态、执行 KV connector bookkeeping，total_num_scheduled_tokens 可能为 0。没有 token 要 forward 时，就不应该做 PP receive/send。

intermediate_tensors 表示从上一个 PP stage 收到的中间激活。第一个 PP stage 没有上游，所以保持 None；非第一个 PP stage 会在后面通过 irecv_tensor_dict(...) 得到它。

all_gather_tensors 则是 PP 通信时给 communicator 的额外提示。它不表示 layer 内部的 TP all-reduce，而是表示跨 PP stage 传 tensor 时，某些 tensor 是否需要配合 TP group 做 all-gather。

4. PP + SP 下的 all_gather_tensors

这一段只在同时满足三个条件时触发：

if (
    parallel_config.pipeline_parallel_size > 1
    and compilation_config.pass_config.enable_sp
    and forward_pass
):

也就是说，只有在开启 pipeline parallelism、开启 sequence parallelism，并且本轮确实有 token 要 forward 时，才需要提前计算 all_gather_tensors。

内部先把每个 request 本轮调度的 token 数量取出来：

num_scheduled_tokens_np = np.array(
    list(scheduler_output.num_scheduled_tokens.values()),
    dtype=np.int32,
)

然后提前调用一次：

1	`self.model_runner._determine_batch_execution_and_padding(...)`

这一步不是正式 forward，而是为了提前得到 batch_desc，尤其是 padding 后的 batch_desc.num_tokens。sequence parallel 下，某些 residual tensor 是否已经 scatter，和实际 batch 形态、padding 后 token 数有关。

最终得到：

all_gather_tensors = {
    "residual": not is_residual_scattered_for_sp(
        self.vllm_config, batch_desc.num_tokens
    )
}

它的含义是：跨 PP stage 传递 residual 时，如果 residual 当前不是 sequence-parallel scatter 状态，那么 PP communicator 需要借助 TP group 做 all-gather，让下游 stage 拿到它期望的 tensor 形态。

这不是 Transformer layer 内部的 TP all-reduce。它发生在 PP stage 边界，用来保证跨 stage 传递的 tensor 形态正确。

5. 非第一个 PP stage 接收上游激活

如果当前 rank 不属于第一个 PP stage，并且本轮确实有 forward，那么它需要从上一个 PP stage 接收 intermediate tensors：

if forward_pass and not get_pp_group().is_first_rank:
    tensor_dict, comm_handles, comm_postprocess = (
        get_pp_group().irecv_tensor_dict(
            all_gather_group=get_tp_group(),
            all_gather_tensors=all_gather_tensors,
        )
    )

这段代码表达了 PP 的基本数据流：第一个 PP stage 从 token/input embedding 开始算；后续 PP stage 不再从原始 token 开始，而是从上一个 stage 的 hidden states 继续算。

irecv_tensor_dict(...) 返回三类对象：

tensor_dict:
  上游 stage 发来的 tensor 字典。

comm_handles:
  非阻塞通信句柄。

comm_postprocess:
  通信完成后需要执行的后处理，可能包含 TP all-gather 相关逻辑。

随后它们被包装成：

intermediate_tensors = AsyncIntermediateTensors(
    tensor_dict,
    comm_handles=comm_handles,
    comm_postprocess=comm_postprocess,
)

这层包装的意义是让通信和本地准备工作重叠。当前 stage 可以先发起 receive，然后在 GPUModelRunner 里做 request state 更新、metadata 构造等工作；等真正需要上游 hidden states 时，再同步通信句柄并执行 postprocess。

6. 进入 GPUModelRunner

完成 PP 输入准备后，GPUWorker 把执行交给 GPUModelRunner：

with self.annotate_profile(scheduler_output):
    output = self.model_runner.execute_model(
        scheduler_output, intermediate_tensors
    )

这是整段函数的核心边界。GPUWorker 只负责传入 scheduler_output 和可能存在的 intermediate_tensors；真正复杂的推理执行在 GPUModelRunner.execute_model() 内部完成，包括：

_update_states:
  根据 SchedulerOutput 更新本地 request cache 和 persistent batch。

_prepare_inputs:
  构造 input_ids、positions、query_start_loc、seq_lens 等输入张量。

_get_slot_mappings:
  根据 block table 生成 attention 写 KV cache 时需要的 slot mapping。

_build_attention_metadata:
  构造 attention backend 所需的 block table、seq lens、slot mapping 等 metadata。

_preprocess:
  决定使用 input_ids、inputs_embeds，或 PP intermediate_tensors。

_model_forward:
  真正调用 self.model(...)。

如果是第一个 PP stage，intermediate_tensors 为 None，模型从 token/input embeddings 开始；如果是后续 PP stage，GPUModelRunner 会把 intermediate_tensors 作为本 stage forward 的输入。

7. V2 pooling 的特殊处理

这段是 V2 ModelRunner 的 pooling 模型特例：

if (
    self.use_v2_model_runner
    and self.model_runner.is_pooling_model
    and output is None
):
    output = self.model_runner.pool()

普通生成模型可以先忽略它。它的意义是：某些 pooling 模型在 execute_model 后不会直接返回最终结果，需要额外调用 pool() 才能得到 pooling output。

这段逻辑放在 GPUWorker 顶层，是为了让 executor 看到的返回值仍然符合统一接口。

8. 根据返回值决定执行路径

model_runner.execute_model(...) 的返回值决定 GPUWorker 接下来做什么：

if isinstance(
    output, ModelRunnerOutput | AsyncModelRunnerOutput | NoneType
):
    return output

这里有三类“可以直接返回”的结果：

ModelRunnerOutput:
  已经得到最终输出，可以返回给 executor。

AsyncModelRunnerOutput:
  输出存在异步 GPU -> CPU copy，executor 后续会等待或消费。

None:
  当前 rank 不产生最终输出，或者 execute_model 阶段只完成了 forward，
  后续还会通过 sample_tokens 完成采样。

None 不是“没执行”。在 vLLM 的 async scheduling / sampling 设计里，execute_model 可能先完成 forward，把 logits 和必要状态暂存在 model runner 内部；随后 EngineCore 再广播 sample_tokens(grammar_output) 完成 structured output mask 和 sampling。

如果返回值不是这三类，就必须是：

1	`assert isinstance(output, IntermediateTensors)`

这说明当前 stage 不是最后一个 PP stage，它完成了自己负责的 layers，并产出了要传给下一个 PP stage 的 intermediate tensors。

9. 非最后 PP stage 发送中间激活

当 output 是 IntermediateTensors 时，当前 rank 需要把它发给下一个 PP stage：

self._pp_send_work = get_pp_group().isend_tensor_dict(
    output.tensors,
    all_gather_group=get_tp_group(),
    all_gather_tensors=all_gather_tensors,
)

return None

这里同样是 non-blocking send。函数不在这里等待通信完成，而是把通信句柄保存到 _pp_send_work，留到下一轮 execute_model 开头统一 wait()。

这也是为什么函数开头要先等待上一轮 _pp_send_work。PP send 和下一轮计算之间形成了一个轻量 pipeline：尽量让通信异步推进，但在可能复用 buffer 之前保证通信完成。

这段前面还有一个断言：

assert (
    parallel_config.distributed_executor_backend != "external_launcher"
    and not get_pp_group().is_last_rank
)

它表达的是：最后一个 PP stage 不应该返回 IntermediateTensors，因为它已经没有下游 stage；只有非最后 stage 才能走 isend_tensor_dict(...)。

10. 不同并行配置下的行为

GPUWorker.execute_model() 的行为会随着 PP 配置改变。最简单的情况是 PP=1。此时没有上游也没有下游，函数不会做 PP receive/send，基本就是直接调用 model_runner.execute_model(...) 并返回结果。

PP=1:
  no irecv
  model_runner.execute_model(...)
  no isend
  return ModelRunnerOutput / AsyncModelRunnerOutput / None

如果 PP=2，两个 stage 的行为不同：

PP stage 0:
  不接收 intermediate_tensors
  执行前半段 layers
  返回 IntermediateTensors
  isend 给 PP stage 1
  return None

PP stage 1:
  irecv PP stage 0 发来的 intermediate_tensors
  执行后半段 layers
  计算 logits 或准备 sampling state
  return ModelRunnerOutput / AsyncModelRunnerOutput / None

如果再叠加 TP=8，每个 PP stage 内部有 8 个 TP ranks。跨 PP stage 的 tensor 通信通常沿 TP lane 进行：

PP0.TP0 -> PP1.TP0
PP0.TP1 -> PP1.TP1
...
PP0.TP7 -> PP1.TP7

这使得 PP stage 间的数据传输不会集中在单个 rank 上，而是由 TP group 内多个 ranks 分担。

11. TP 内部的基本不变量

在 Megatron-style tensor parallelism 中，一个 TP group 内的多个 rank 共同执行同一个 layer。每个 rank 只持有部分权重，局部计算完成后，通过 collective 让后续计算拿到正确形态的 hidden states。

对于 transformer block 来说，残差流上的主张量通常可以抽象成：

1	`[N, hidden]`

这里的 N 是本轮实际参与 forward 的 token 数，可以来自 prefill、decode 或混合 batch；hidden 是模型隐藏维度。

在很多关键边界上，TP group 内每个 rank 都需要看到完整的 [N, hidden] 残差状态。这个完整状态通常不是每个 rank 独立算出来的，而是通过 TP collective 从局部结果合成出来的。

12. Column Parallel 和 Row Parallel 的配合

TP 的核心技巧不是“每一层末尾都无条件 all-reduce”，而是让 column-parallel 和 row-parallel 成对配合，尽量减少通信次数。

一个典型线性层可以写成：

Y = X A

Column parallel 是按输出维度切权重：

1 2	`A = [A0, A1, ..., A(tp-1)] Yi = X Ai`

每个 rank 得到的是输出 hidden 的一部分，所以局部结果形状更像：

1	`[N, hidden / TP]`

Row parallel 是按输入维度切权重：

X = [X0, X1, ..., X(tp-1)]
A = [A0; A1; ...; A(tp-1)]
Yi = Xi Ai
Y = sum_i Yi

每个 rank 先算出一个 partial output，然后通过 all_reduce 求和。all_reduce 之后，每个 TP rank 都拿到完整的：

1	`[N, hidden]`

这就是“每个 TP rank 最后都有完整 hidden”的来源。更精确地说，它通常发生在 row-parallel 的输出边界，而不是每一个 column-parallel projection 后都发生。

13. Attention 中的 TP 通信

Attention 的 QKV projection 通常使用 QKVParallelLinear，它本质上是 column-parallel。每个 TP rank 负责一部分 attention heads。

可以把 attention 的 TP 数据流理解成：

输入 residual:
  每个 rank 都有完整 [N, hidden]

QKV projection:
  每个 rank 计算自己负责的 Q/K/V heads
  局部结果是部分 heads

Attention kernel:
  每个 rank 在自己的 heads 上做 attention
  写自己负责 heads 的 KV cache

O projection:
  RowParallelLinear
  每个 rank 产生 partial [N, hidden]
  all_reduce 后每个 rank 得到完整 [N, hidden]

画成图：

flowchart TD
    X["replicated X: [N, hidden]"] --> QKV["QKVParallelLinear<br/>column parallel"]
    QKV --> Heads["local Q/K/V heads<br/>per TP rank"]
    Heads --> Attn["attention per local heads"]
    Attn --> O["O projection<br/>RowParallelLinear"]
    O --> AR["all_reduce across TP group"]
    AR --> Y["replicated Y: [N, hidden]"]

所以，attention block 的关键通信通常发生在 output projection 后。QKV projection 本身把 heads 分给不同 TP rank，attention kernel 在本 rank 的 heads 上执行；最后通过 row-parallel output projection 的 all_reduce 回到完整 hidden。

14. MLP 中的 TP 通信

MLP 里也有类似结构。以常见 gated MLP 为例，gate_proj 和 up_proj 通常合并成 MergedColumnParallelLinear，down_proj 通常是 RowParallelLinear。

数据流可以理解为：

输入 residual:
  每个 rank 都有完整 [N, hidden]

gate/up projection:
  ColumnParallel / MergedColumnParallel
  每个 rank 计算 intermediate hidden 的一部分

activation:
  在局部分片上完成

down projection:
  RowParallelLinear
  每个 rank 计算 partial [N, hidden]
  all_reduce 后每个 rank 得到完整 [N, hidden]

因此，MLP 的通信边界通常也在 down_proj 后。column-parallel 扩展阶段保持分片，row-parallel 收缩阶段通过 all_reduce 回到完整 residual。

15. PP 边界为什么还需要 TP all_gather

GPUWorker.execute_model() 里的 PP 通信不是简单地“某个 rank 把完整 tensor 发给下一个 stage”。在 TP=8 时，每个 PP stage 内都有 8 个 TP rank；跨 PP stage 传 tensor 时，通常是对应 TP lane 之间传：

PP0 TP0 -> PP1 TP0
PP0 TP1 -> PP1 TP1
...
PP0 TP7 -> PP1 TP7

这样做可以避免单个 rank 负责全部跨 stage 数据传输。也就是说，每个 TP rank 负责传输 X[i] 的一部分，从而让跨 PP stage 的带宽压力分散到整个 TP group。

不过，下一个 PP stage 期望收到的 tensor 形态必须和它的计算方式一致。如果当前 tensor 在 TP group 内是分片状态，那么可以按 lane 传分片；如果下游需要完整 tensor，或者某些 tensor 没有按 sequence parallel 的方式 scatter，就需要在 TP group 内先做 all_gather。

这也是 GPUWorker.execute_model() 中 all_gather_tensors 的含义：

all_gather_tensors = {
    "residual": not is_residual_scattered_for_sp(
        self.vllm_config, batch_desc.num_tokens
    )
}

它不是在做模型 layer 的 TP all-reduce，而是在告诉 PP communicator：跨 PP stage 传递某些 tensor 时，是否需要在 TP group 内额外做 all-gather。

16. 总结

GPUWorker.execute_model() 是具体推理链路的第一个关键入口。它不直接解释 Transformer layer 内部如何计算，而是负责把 SchedulerOutput 带进 GPUModelRunner，并在 PP stage 之间正确传递 intermediate tensors。

这段代码可以压缩成三层逻辑：

PP 边界:
  非 first stage 接收上游 intermediate tensors；
  非 last stage 发送下游 intermediate tensors。

ModelRunner 执行:
  更新本地 request state；
  构造 input tensors 和 attention metadata；
  调用 self.model(...) 完成本 stage forward。

TP 辅助:
  layer 内部用 all_reduce 等 collective 合成正确 hidden；
  PP 边界必要时用 all_gather_tensors 决定跨 stage 传输形态。

因此，这段函数里的通信可以分成三类：PP stage 间的 tensor 传递、TP layer 内部的计算同步，以及 PP 边界上为了 tensor 形态正确而触发的 TP-assisted all-gather。

#vLLM #code-reading #LLM serving #systems

vLLM Code Reading (Section 5): GPUWorker.execute_model 与 PP/TP 通信

https://jeremyguo.space/2026/06/26/vllm-code-reading-section-5-gpuworker-pp-tp-communication/

作者

郭俊毅 / JeremyGuo

发布于

2026年6月26日

许可协议

vLLM Code Reading (Section 6): KV Cache 中的 BlockTable 与 Slot Mapping 上一篇

vLLM Code Reading (Section 4): EngineCore 到 Worker 的执行链路下一篇