Prhub

#43732 [Core] Cleanup KVConnector handling with PP + fix MRV2

原始 PR 作者 njhill 合并时间 2026-05-29 04:12 文件变更 9 提交数 6 评论 3 代码增减 +38 / -48

执行摘要

清理 KVConnector 处理并修复 PP 下 MRV2 输出

根据 PR 描述,此举旨在清理不再使用的 IntermediateTensors.kv_connector_output 字段,并简化各处需要返回仅包含 kv connector output 的空 ModelRunnerOutput 的逻辑。更重要的是修复 MRV2 非最终 PP rank 应输出 kv connector output 用于聚合的问题,确保跨节点 KV 传输的正确性。

值得精读,特别是 with_kv_conn_output_only 静态方法的设计:通过一个工厂方法统一了空输出创建的逻辑,避免了多处重复的 copy 和判断。这种模式在需要返回带可选字段的空容器时非常有用,可以推广到项目中其他类似场景。

讨论亮点

WoosukKwon 在 review 时指出 tests/v1/worker/test_gpu_model_runner.py 中对于 sample_tokens 返回值的断言应使用 in 而非 is,以同时兼容 EMPTY_MODEL_RUNNER_OUTPUTNone 的情况。njhill 采纳并修改了测试。

实现拆解

  1. vllm/v1/outputs.py 中为 ModelRunnerOutput 新增静态方法 with_kv_conn_output_only,该方法接受 KVConnectorOutput 参数,若为 None 或空则返回全局 EMPTY_MODEL_RUNNER_OUTPUT,否则通过浅拷贝空实例并设置 kv_connector_output 字段返回。
  2. vllm/sequence.py 中从 IntermediateTensors 类中移除 kv_connector_output 字段及其相关类型导入(KVConnectorOutput, TYPE_CHECKING, Any),因为各模型运行器现在自行管理该状态,不再需要附加在中间张量上传递。
  3. vllm/v1/worker/kv_connector_model_runner_mixin.py 中,将 kv_connector_no_forward 方法中的手动检查与构造代码精简为调用 ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output),同时移除了不再需要的 import copyEMPTY_MODEL_RUNNER_OUTPUT 导入引用。
  4. vllm/v1/worker/gpu_model_runner.pyvllm/v1/worker/gpu/kv_connector.py 中执行相似的简化,替换重复的空输出创建逻辑。
  5. vllm/v1/worker/gpu/model_runner.py(即 MRV2)中,修复 sample_tokenspool 方法中非最终 PP rank 的返回路径:之前对非最终 rank 直接返回 NoneEMPTY_MODEL_RUNNER_OUTPUT,现改为通过 ModelRunnerOutput.with_kv_conn_output_only 返回包含 KV connector output 的输出,确保聚合层能正确收到跨节点传输结果。
  6. 对应调整测试文件 tests/v1/worker/test_gpu_model_runner.pytests/v1/worker/test_gpu_model_runner_v2_eplb.py 中的断言,使其兼容新的输出类型(可能为空输出或 None)。
文件 模块 状态 重要度
vllm/v1/outputs.py 输出层 modified 6.99
vllm/v1/worker/kv_connector_model_runner_mixin.py KV 连接器 modified 6.26
vllm/v1/worker/gpu_model_runner.py 运行器 v1 modified 6.15
vllm/sequence.py 序列 modified 5.71
vllm/v1/worker/gpu/model_runner.py MRV2 modified 6.39

关键符号

with_kv_conn_output_only

关键源码片段

vllm/v1/outputs.py core-logic

核心变更:新增 with_kv_conn_output_only 静态方法,作为统一入口创建仅包含 KV connector output 的 ModelRunnerOutput 实例。

from copy import copy
from dataclasses import dataclass, field
# ... 其他导入@dataclass
class ModelRunnerOutput:
    # ... 其他字段
    kv_connector_output: KVConnectorOutput | None = None
    # ...
​
    @staticmethod
    def with_kv_conn_output_only(
        kv_connector_output: KVConnectorOutput | None,
    ) -> "ModelRunnerOutput":
        """Return ModelRunnerOutput containing the provided KVConnectorOutput,
        otherwise empty. Returns EMPTY_MODEL_RUNNER_OUTPUT if
        kv_connector_output is None or empty.
        """
        # 如果输入为 None 或空,则返回全局单例空输出
        if kv_connector_output is None or kv_connector_output.is_empty():
            return EMPTY_MODEL_RUNNER_OUTPUT
        # 浅拷贝空实例,仅填充 kv_connector_output 字段
        output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
        output.kv_connector_output = kv_connector_output
        return output
vllm/sequence.py dependency-wiring

从 IntermediateTensors 中移除了 kv_connector_output 字段和相关类型导入,简化数据结构。

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM projectfrom dataclasses import dataclass
import torch@dataclass
class IntermediateTensors:
    """For all pipeline stages except the last, we need to return the hidden
    states and residuals to be sent to the next stage. This data structure
    contains the hidden states and residuals for a request.
    """
    tensors: dict[str, torch.Tensor]
​
    def __init__(self, tensors: dict[str, torch.Tensor]) -> None:
        # 手动定义 init 以确保 Dynamo 能够追踪来源文件
        self.tensors = tensors
​
    def __getitem__(self, key: str | slice):
        if isinstance(key, str):
            return self.tensors[key]
        elif isinstance(key, slice):
            return self.__class__(
                {k: v[key] for k, v in self.tensors.items()}
            )
​
    def __setitem__(self, key: str, value: torch.Tensor):
        self.tensors[key] = value
​
    # items, __len__, __eq__, __repr__ 等方法保持不变

评论区精华

测试断言格式优化 style

WoosukKwon 指出测试中 assert 表达式应使用 in 而非 is,以便更清晰地表达可能返回两种值(EMPTY_MODEL_RUNNER_OUTPUT 或 None)。

结论:njhill 接受并修改为 assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)。 · 已解决

风险与影响

主要风险在于移除了 IntermediateTensors.kv_connector_output 字段,若其他代码路径仍尝试访问该属性将导致 AttributeError。通过搜索仓库发现该字段仅在本次涉及的文件中使用,风险可控。新的静态方法在所有调用点被替换,但需注意在 gpu_model_runner.pyexecute_model 的非最后 rank 分支清理了 hidden_states.kv_connector_output = kv_connector_output 赋值,若下游仍有依赖该属性的代码可能受影响。此外,MRV2 修复改变了非最后 rank 的返回类型,可能影响调用者,但测试已做适配。整体回归风险较低。

对用户透明,无 API 变化。内部影响:修复了启用 PP 和 KV connector 时 MRV2 场景下跨节点 KV 传输的正确性,可能改善基于 MRV2 的服务的稳定性和准确性。代码结构更清晰,维护成本降低。团队后续在类似场景可复用 with_kv_conn_output_only 方法。

字段移除可能遗漏 测试覆盖有限 核心路径变更

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论