执行摘要

• 一句话：为推测解码添加运行时草稿权重更新
• 推荐动作：此 PR 实现了推测解码场景中缺失的 draft 权重运行时更新能力，设计上通过 set_weight_update_target 实现了 WeightTransferEngine 的目标切换，并妥善处理了 sleep/wake_up 参数保留问题。推荐精读 gpu_worker.py 中的 _start_weight_update 实现和 base.py 中的目标切换方法，它们是整个功能的骨架。同时关注与 #47357 的后续兼容调整。

功能与动机

在 RL 训练循环中，当训练器通过标准路径更新 verifier 权重时，draft 模型的独立参数从未被更新，共享引用可能被静默切断，且 level-2 睡眠会丢弃 draft 参数。此 PR 旨在填补这一缺口，使推测解码的 draft 模型能与 verifier 一起保持最新，从而提升 RL 训练的吞吐和奖励信号准确性。（来源于 PR Body）

实现拆解

实现分为以下步骤：

1. 扩展 WeightTransferEngine 基类：在 vllm/distributed/weight_transfer/base.py 的 WeightTransferEngine 中添加 set_weight_update_target(model, model_config) 和 reset_weight_update_target() 方法，并声明类属性 supports_draft_weight_update = True，使引擎能够临时切换到 draft 模型作为权重更新目标。set_weight_update_target 将当前 model 和 model_config 替换为传入的 draft 模型及其配置；reset_weight_update_target 恢复为构造函数中存储的默认模型。
2. 在 Worker 中实现 draft 模型访问与会话管理：在 vllm/v1/worker/gpu_worker.py 的 Worker 类中新增 get_draft_model() 代理方法和 _set_draft_weight_update_target() 内部方法。新增 start_draft_weight_update() 公共 API，它会调用 _start_weight_update(is_draft=True)，在引擎上调用 set_weight_update_target，并标记 _weight_update_active。finish_weight_update 中加入对 draft 更新会话的检测，调用 reset_weight_update_target。如果引擎不支持 draft 更新（如 SparseNCCLEngine），则抛出 RuntimeError。
3. 保护 sleep/wake_up 中的 draft 参数：在 sleep(level=2) 中，通过 get_draft_model 检测 draft 是否包含 _build_fused_kv_buffers 方法并将标志 _sleep_rebuild_draft_metadata_buffers 设为 True。在 wake_up 中，如果此标志为真，则重建 draft 的融合 KV 缓冲区。同时添加 _sleep_saved_draft_params 字典，确保 draft 参数在 GPU 内存释放后恢复。
4. 暴露协议与 HTTP 端点：在 vllm/engine/protocol.py 的 EngineClient 协议中添加抽象方法 start_draft_weight_update。在 vllm/entrypoints/serve/dev/rlhf/api_router.py 中注册 POST /start_draft_weight_update 路由，调用引擎客户端的方法并返回确认消息。同步更新 vllm/entrypoints/llm.py 和 vllm/v1/engine/async_llm.py 以支持该调用。
5. 配套测试与文档：在 tests/v1/worker/test_gpu_worker_weight_transfer.py 中添加对 reset_weight_update_target 的测试。在 docs/training/weight_transfer/base.md 中添加相关文档。同时废弃原提议的 update_speculative_model_weights 方法，仅保留公共 API 的简洁性。

关键文件：

• vllm/v1/worker/gpu_worker.py（模块 Worker；类别 source；类型 core-logic；符号 get_draft_model, _set_draft_weight_update_target, start_draft_weight_update, _start_weight_update）: 核心逻辑所在，新增 get_draft_model、_set_draft_weight_update_target、start_draft_weight_update、_start_weight_update；修改 sleep/wake_up 以保留 draft 参数。
• vllm/distributed/weight_transfer/base.py（模块 权重传输；类别 source；类型 core-logic；符号 set_weight_update_target, reset_weight_update_target）: WeightTransferEngine 基类扩展了目标切换能力，并声明 draft 支持属性。
• vllm/v1/worker/gpu_model_runner.py（模块 Model Runner；类别 source；类型 data-contract；符号 get_draft_model）: V1 Model Runner 添加 get_draft_model()，通过 drafter.model 访问 draft 模型并处理 wrapper。
• vllm/v1/worker/gpu/model_runner.py（模块 Model Runner；类别 source；类型 data-contract；符号 get_draft_model）: V2 Model Runner 添加 get_draft_model()，通过 speculator.model 访问。
• vllm/engine/protocol.py（模块 引擎协议；类别 source；类型 core-logic；符号 start_draft_weight_update）: EngineClient 协议新增 start_draft_weight_update 抽象方法，确保所有引擎客户端实现该功能。
• vllm/entrypoints/serve/dev/rlhf/api_router.py（模块 API 路由；类别 source；类型 entrypoint；符号 start_draft_weight_update）: RLHF HTTP 路由器注册 POST /start_draft_weight_update 端点，供外部训练框架调用。
• tests/v1/worker/test_gpu_worker_weight_transfer.py（模块 测试；类别 test；类型 test-coverage；符号 reset_weight_update_target）: 测试 reset_weight_update_target 功能，验证目标切换后的恢复。
• docs/training/weight_transfer/base.md（模块 文档；类别 docs；类型 documentation）: 用户文档中增加对 draft 权重更新用法的描述。

关键符号：get_draft_model, _set_draft_weight_update_target, start_draft_weight_update, _start_weight_update, set_weight_update_target, reset_weight_update_target

关键源码片段

vllm/v1/worker/gpu_worker.py

核心逻辑所在，新增 get_draft_model、_set_draft_weight_update_target、start_draft_weight_update、_start_weight_update；修改 sleep/wake_up 以保留 draft 参数。

# vllm/v1/worker/gpu_worker.py ( 部分实现 )
# 启动权重更新，支持切换到 draft 模型
​
def _start_weight_update(self, is_draft: bool = False) -> None:
    self._check_weight_transfer_engine()
    assert self.weight_transfer_engine is not None
​
    # 检查引擎是否支持 draft 更新
    if is_draft and not self.weight_transfer_engine.supports_draft_weight_update:
        raise RuntimeError(
            f"{type(self.weight_transfer_engine).__name__} does not support "
            "draft model weight updates."
        )
​
    # 如果是 draft 更新，将引擎目标切换到 draft 模型
    if is_draft:
        self._set_draft_weight_update_target()
​
    # 委托给引擎执行具体准备
    self.weight_transfer_engine.start_weight_update()
    self._weight_update_active = True
​
​
def finish_weight_update(self) -> None:
    """结束权重更新会话，恢复引擎目标。"""
    if not self._weight_update_active:
        raise RuntimeError(
            "finish_weight_update called without a matching start_weight_update."
        )
    # ... 引擎结束逻辑 ...
    # 如果是 draft 会话，恢复引擎目标到 verifier
    if self._draft_weight_update_active:
        self.weight_transfer_engine.reset_weight_update_target()
        self._draft_weight_update_active = False
    self._weight_update_active = False
​
​
# sleep/wake_up 中保护 draft 融合 KV 缓冲区
​
def sleep(self, level: int = 1) -> None:
    # ...
    if level == 2:
        # 保存 verifier 缓冲区
        model = self.model_runner.model
        self._sleep_saved_buffers = {
            name: buffer.cpu().clone() for name, buffer in model.named_buffers()
        }
        # 检测 draft 是否需要重建融合 KV 缓冲区
        draft = self.get_draft_model()
        inner = getattr(draft, "model", None) if draft is not None else None
        self._sleep_rebuild_draft_metadata_buffers = inner is not None and hasattr(
            inner, "_build_fused_kv_buffers"
        )
    self._get_sleep_mode_backend().suspend(level)
​
​
def wake_up(self, tags: list[str] | None = None) -> None:
    self._get_sleep_mode_backend().resume(tags)
    # 恢复 verifier 缓冲区
    if len(self._sleep_saved_buffers):
        model = self.model_runner.model
        for name, buffer in model.named_buffers():
            if name in self._sleep_saved_buffers:
                buffer.data.copy_(self._sleep_saved_buffers[name].data)
        self._sleep_saved_buffers = {}
    # 重建 draft 融合 KV 缓冲区
    if self._sleep_rebuild_draft_metadata_buffers:
        draft = self.get_draft_model()
        if draft is not None:
            inner = getattr(draft, "model", None)
            if inner is not None and hasattr(inner, "_build_fused_kv_buffers"):
                inner._build_fused_kv_buffers()
        self._sleep_rebuild_draft_metadata_buffers = False

vllm/distributed/weight_transfer/base.py

WeightTransferEngine 基类扩展了目标切换能力，并声明 draft 支持属性。

# vllm/distributed/weight_transfer/base.py ( 部分实现 )
​
class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
    # 类属性，表示该引擎是否支持 draft 权重更新
    supports_draft_weight_update: bool = True
​
    def __init__(
        self,
        config: WeightTransferConfig,
        vllm_config: "VllmConfig",
        device: torch.device,
        model: torch.nn.Module,
    ) -> None:
        # ...
        self.model = model
        self.model_config = vllm_config.model_config
        # 保存默认目标，用于重置
        self._default_model_config = self.model_config
        self._default_model = model
​
    def set_weight_update_target(
        self,
        model: torch.nn.Module,
        model_config: Any,
    ) -> None:
        """将当前活动的权重更新目标设置为给定模型。"""
        self.model = model
        self.model_config = model_config
​
    def reset_weight_update_target(self) -> None:
        """将权重更新目标恢复为引擎的默认模型。"""
        self.model = self._default_model
        self.model_config = self._default_model_config

vllm/v1/worker/gpu_model_runner.py

V1 Model Runner 添加 get_draft_model()，通过 drafter.model 访问 draft 模型并处理 wrapper。

# vllm/v1/worker/gpu_model_runner.py ( 部分实现 )
​
def get_draft_model(self) -> nn.Module | None:
    drafter = getattr(self, "drafter", None)
    if drafter is None:
        return None
    model = getattr(drafter, "model", None)
    # 如果被 CUDAGraph 等 wrapper 包裹，则解包
    if isinstance(
        model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper)
    ):
        return cast(nn.Module, model.unwrap())
    return cast(nn.Module | None, model)

评论区精华

Review 中主要讨论了以下核心问题：

• 死代码争议：aoshen02 指出 update_speculative_model_weights 是死代码，vx120 解释这是计划给外部 RL 框架（如 verl）使用的公共 API，并非 vLLM 内部使用。最终该代码被移除，统一使用 start_draft_weight_update/finish_weight_update 生命周期。
• 接口设计分歧：aoshen02 建议在 start_weight_update 中添加 include_draft 标志来复用同一端点，但后来采纳了独立方法 start_draft_weight_update，以保持协议清晰且避免污染 verifier 的更新流程。
• get_draft_model 实现简化：TheEpicDolphin 建议在 V2 Model Runner 中直接使用 isinstance(self.speculator, DraftModelSpeculator) 判断，并在 V1 中直接通过 drafter.model 访问，避免调用 get_model() 这样的重方法。作者采纳并精简了实现。
• 日志记录器分离：aoshen02 要求将 _runtime_stat_logger_factories 等指标相关代码移出此 PR，作者同意并删除，承诺单独提交。

• Sleep 模式参数保留：aoshen02 指出 level-2 sleep 也需要保留 draft 的参数，代码中通过保存 draft 命名缓冲区并在 wake_up 中重建融合 KV 缓冲区来处理。

  • update_speculative_model_weights 死代码争议 (design): 该代码被移除，统一使用 start_draft_weight_update/finish_weight_update 生命周期。
  • 接口设计分歧：include_draft 标志 vs 独立方法 (design): 采用独立方法 start_draft_weight_update，团队达成一致。
  • get_draft_model 实现简化 (design): 简化后的实现被合并。
  • 日志记录器分离到独立 PR (design): 代码被移除，承诺单独 PR。
  • Sleep level-2 时 draft 参数保留 (correctness): 已实现并合入，通过 _sleep_rebuild_draft_metadata_buffers 标志控制。

风险与影响

• 风险：兼容性风险：新 start_draft_weight_update 协议方法要求所有 EngineClient 子类实现；若第三方自定义引擎未实现该方法，在调用时会抛出 NotImplementedError。
  睡眠恢复完整性：level-2 sleep 的 draft 参数恢复仅针对带有 _build_fused_kv_buffers 的 DFlash 模型，其他 draft 架构可能无法正确恢复，需未来扩展。
  Sparse NCCL 引擎限制：SparseNCCLEngine 明确拒绝 draft 权重更新，若在 RL 训练中使用该引擎且尝试更新 draft，会抛出 ValueError，可能影响用户的工作流。
  来自 #47357 的潜在冲突：hao-aaron 在最终批准中指出 #47357 正在引入 stateful trainer 对象，其中需要指定哪个模型接收更新，与本 PR 的接口可能需要协调。

• 影响：用户影响：RL 训练用户现在可以在推测解码场景下更新 draft 模型权重，显著提升长序列生成质量和奖励准确性。需要配合 RL 框架（如 verl）调用新端点 POST /start_draft_weight_update。
  系统影响：修改了 Worker、Engine、HTTP 三层，但向后兼容：已有 start_weight_update/update_weights/finish_weight_update 路径未改变。draft 更新路径完全独立。
  团队影响：需要确保后续 #47357 的合并不会冲突；未来可能将 draft 权重更新集成到统一的 stateful trainer 接口中。

• 风险标记：核心路径变更, 睡眠模式影响, 兼容性风险, 缺少 E2E 测试覆盖

关联脉络

• PR #47357 [Stateful Trainer] 指定模型更新目标: hao-aaron 在 review 中提及此 PR 正在引入 stateful trainer 对象，与本 PR 的接口可能需要协调，建议先合并本 PR 然后由他适配。
• PR #45586 [Spec Decode] E2E 测试: aoshen02 建议将本 PR 的单元测试移除，改为将端到端测试集成到 #45586 中。
• PR #28257 SGLang draft 权重更新修复: PR body 中引用该 SGLang issue，说明相同问题在其他推理引擎中也存在，提供了跨项目背景。