执行摘要

• 一句话：Breakable CUDA 图仅接管 PIECEWISE 路径，FULL 恢复标准捕获
• 推荐动作：值得精读。两个设计决策尤其值得关注：（1）wrapper 的 runtime_mode 门控如何同时服务 MRV1（严格匹配）与 MRV2（默认 match-any），一份代码两种契约；（2）load_model 中“内层 breakable 先装、外层 FULL 再套”的独立 if 结构，如何在不显式依赖 mode 枚举的情况下吸收 drafter 与后期模式升级两类边界情况。结合 review 评论与 #55095 阅读，可以完整理解 vLLM 当前 CUDA Graph 捕获层的演进脉络。

功能与动机

PR body 明确指出：VLLM_USE_BREAKABLE_CUDAGRAPH=1 no longer overrides FULL cudagraphs in ModelRunnerV1——在 cudagraph_mode=FULL_AND_PIECEWISE 下，decode 分发改用标准 CUDAGraphWrapper(FULL)，breakable 捕获只需替换 PIECEWISE，从而对齐 MRV2 的 use_breakable_cg 门控。此前 breakable 作为实验性路径会整条劫持成熟稳定的 FULL 捕获，既偏离 MRV2 行为，也让 decode 主路径承担了不必要的风险。此外 AndreaasKaratzas 在 review 中强调要避免 #55095 中“选择 piecewise 却没有捕获者”的回归，本次包裹策略调整正是围绕这一目标展开。

实现拆解

变更入口是 MRV1 模型装配路径 vllm/v1/worker/gpu_model_runner.py::load_model()，配套修改两个 wrapper 实现与单元测试，分四步：

1. 为 BreakableCUDAGraphWrapper 增加 runtime_mode 门控（vllm/compilation/breakable_cudagraph.py）
   - __init__ 新增可选参数 runtime_mode: CUDAGraphMode | None = None 并存入 self.runtime_mode。
   - __call__ 中，当 cudagraph_runtime_mode == NONE，或显式设置了 runtime_mode 且与当前分发模式不匹配时，直接 eager 调用 self.runnable；仅匹配时才走 _capture/_replay。
   - 默认 None 保留旧的 match-any 行为，保证 MRV2 调用点（不传 runtime_mode）零行为变化；同时允许 MRV1 将 breakable 严格限定到 PIECEWISE。

2. 重构 load_model 的包裹顺序（vllm/v1/worker/gpu_model_runner.py）
   - breakable 分支条件保持 cudagraph_mode != NONE（而非中间版本收窄为 has_piecewise_cudagraphs()），但包裹时显式传 runtime_mode=PIECEWISE；drafter.model 同样无条件包裹。
   - 将原先的 elif 改为独立 if：当 cudagraph_mode.has_full_cudagraphs() 时，外层再套 CUDAGraphWrapper(..., runtime_mode=FULL)。
   - 这样在 FULL_AND_PIECEWISE 下形成 CUDAGraphWrapper(FULL) -> BreakableCUDAGraphWrapper(PIECEWISE) 嵌套；在纯 FULL 下，主模型的 FULL 由外层标准 wrapper 捕获，而 drafter（只分发 PIECEWISE 或 NONE）仍由内层 breakable 捕获，避免 drafter eager 运行；后期 FULL -> FULL_AND_PIECEWISE 升级（_check_and_update_cudagraph_mode）也能被已安装的 breakable 覆盖。

3. unwrap() 递归穿透（vllm/compilation/breakable_cudagraph.py 与 vllm/compilation/cuda_graph.py）
   - 两个 wrapper 的 unwrap() 均改为：若 runnable 自身带 unwrap，则递归调用其 unwrap()，否则返回 runnable；保证嵌套后 get_model()/unwrap() 仍能直达原始模型。

4. 单元测试配套（tests/v1/cudagraph/test_breakable_cudagraph.py，+131）
   - 新增 _mock_vllm_config() 与 4 个测试：runtime_mode 不匹配时 fall through（断言 eager 调用且 entries 为空）、匹配时触发 _capture、默认 None 匹配任意非 NONE 模式（注释明确 MRV2 依赖此行为）、嵌套 wrapper 下 unwrap() 直达原始 model。
   - 测试通过 patch.object(BreakableCUDAGraphWrapper, "_capture") 隔离真实捕获，验证的是 dispatch 契约而非 CUDA 行为，可在无 GPU 环境运行。

关键文件：

• vllm/v1/worker/gpu_model_runner.py（模块 模型加载；类别 source；类型 core-logic；符号 load_model）: MRV1 模型装配入口，包裹策略从 elif 改为独立 if，形成 FULL 标准 wrapper 在外、breakable(PIECEWISE) 在内的嵌套结构，是本 PR 行为变更的落点。
• vllm/compilation/breakable_cudagraph.py（模块 编译层；类别 source；类型 core-logic；符号 BreakableCUDAGraphWrapper.init, BreakableCUDAGraphWrapper.call, BreakableCUDAGraphWrapper.unwrap）: BreakableCUDAGraphWrapper 增加 runtime_mode 门控与递归 unwrap，实现 PIECEWISE 严格匹配与 MRV2 match-any 双契约。
• vllm/compilation/cuda_graph.py（模块 编译层；类别 source；类型 core-logic；符号 CUDAGraphWrapper.unwrap）: 标准 CUDAGraphWrapper 的 unwrap() 同步改为递归，保证嵌套后能穿透 breakable 内层直达原始模型。
• tests/v1/cudagraph/test_breakable_cudagraph.py（模块 单元测试；类别 test；类型 test-coverage；符号 _mock_vllm_config, test_wrapper_falls_through_on_runtime_mode_mismatch, test_wrapper_captures_on_runtime_mode_match, test_wrapper_default_matches_any_non_none_mode）: 新增 4 个单元测试锁定 runtime_mode 门控与嵌套 unwrap 契约，是行为回归的主要防线。

关键符号：load_model, BreakableCUDAGraphWrapper.init, BreakableCUDAGraphWrapper.call, BreakableCUDAGraphWrapper.unwrap, CUDAGraphWrapper.unwrap

关键源码片段

vllm/v1/worker/gpu_model_runner.py

MRV1 模型装配入口，包裹策略从 elif 改为独立 if，形成 FULL 标准 wrapper 在外、breakable(PIECEWISE) 在内的嵌套结构，是本 PR 行为变更的落点。

        # wrap the model with full cudagraph wrapper if needed.
        cudagraph_mode = self.compilation_config.cudagraph_mode
        assert cudagraph_mode is not None
        if (
            is_breakable_cudagraph_enabled()
            and cudagraph_mode != CUDAGraphMode.NONE
            and not self.parallel_config.use_ubatching
        ):
            # 限定 PIECEWISE 分发：FULL 由下方标准 wrapper 处理，互不影响。
            # PIECEWISE 分发也可能在包裹之后才出现（例如 FULL target 模式
            # 下的 drafter，或 _check_and_update_cudagraph_mode 后期把
            # FULL 升级为 FULL_AND_PIECEWISE），所以这里无条件先装 breakable。
            self.model = BreakableCUDAGraphWrapper(
                self.model, self.vllm_config, runtime_mode=CUDAGraphMode.PIECEWISE
            )
            drafter = getattr(self, "drafter", None)
            if drafter is not None and hasattr(drafter, "model"):
                # drafter 只会以 PIECEWISE 或 NONE 分发，统一交给 breakable 捕获。
                drafter.model = BreakableCUDAGraphWrapper(
                    drafter.model,
                    self.vllm_config,
                    runtime_mode=CUDAGraphMode.PIECEWISE,
                )
​
        # FULL_AND_PIECEWISE 下外层再套标准 FULL wrapper，形成
        # CUDAGraphWrapper(FULL) -> BreakableCUDAGraphWrapper(PIECEWISE) 的嵌套对。
        if (
            cudagraph_mode.has_full_cudagraphs()
            and not self.parallel_config.use_ubatching
        ):
            self.model = CUDAGraphWrapper(
                self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
            )
        elif self.parallel_config.use_ubatching:
            if cudagraph_mode.has_full_cudagraphs():
                self.model = UBatchWrapper(
                    self.model, self.vllm_config, CUDAGraphMode.FULL, self.device
                )
            else:
                self.model = UBatchWrapper(
                    self.model, self.vllm_config, CUDAGraphMode.NONE, self.device
                )

vllm/compilation/breakable_cudagraph.py

BreakableCUDAGraphWrapper 增加 runtime_mode 门控与递归 unwrap，实现 PIECEWISE 严格匹配与 MRV2 match-any 双契约。

class BreakableCUDAGraphWrapper:
    """Drop-in replacement for :class:`CUDAGraphWrapper`，使用可分段捕获。
​
    与 CUDAGraphWrapper 相同的分发契约：
      * 无 forward_context 时 eager 执行；
      * 运行时模式为 NONE 或与 runtime_mode 不匹配时 eager 执行；
      * 否则按 batch_descriptor 惰性捕获并按需回放。
​
    设置 runtime_mode 后，其他非 NONE 模式会 fall through 到 runnable，
    从而允许外层再套一个处理其他模式的 CUDAGraphWrapper（如 FULL）。
    """
​
    def __init__(
        self,
        runnable: Callable[..., Any],
        vllm_config: VllmConfig,
        runtime_mode: CUDAGraphMode | None = None,
    ) -> None:
        # runtime_mode=None 表示匹配任意非 NONE 分发模式（MRV2 依赖此默认值）；
        # 条目仍按 BatchDescriptor 建键，其本身已编码 batch 形状与均匀性。
        self.runnable = runnable
        self.vllm_config = vllm_config
        self.runtime_mode = runtime_mode
        self.compilation_config = vllm_config.compilation_config
        self.graph_pool = current_platform.get_global_graph_pool()
        self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
        self.entries: dict[BatchDescriptor, _BreakableEntry] = {}
        BreakableCUDAGraphWrapper._all_instances.add(self)
        logger.info_once("Breakable CUDA graph enabled")
​
    def unwrap(self) -> Callable[..., Any]:
        # 递归穿透嵌套 wrapper，直达最底层的原始 runnable。
        runnable = self.runnable
        return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
​
    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if not is_forward_context_available():
            return self.runnable(*args, **kwargs)
​
        forward_context = get_forward_context()
        batch_descriptor = forward_context.batch_descriptor
        cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
​
        # 默认（runtime_mode=None）匹配任意非 NONE 模式；一旦显式设置
        # runtime_mode，其他模式由外层 wrapper（如 FULL 的 CUDAGraphWrapper）
        # 负责捕获，这里直接 eager 调用。
        if cudagraph_runtime_mode == CUDAGraphMode.NONE or (
            self.runtime_mode is not None
            and cudagraph_runtime_mode != self.runtime_mode
        ):
            return self.runnable(*args, **kwargs)
​
        assert batch_descriptor is not None
        entry = self.entries.get(batch_descriptor)
        if entry is None:
            entry = _BreakableEntry(batch_descriptor=batch_descriptor)
            self.entries[batch_descriptor] = entry
​
        if entry.capture is None:
            return self._capture(entry, args, kwargs)
        return self._replay(entry, args, kwargs)

评论区精华

两条 review 评论均来自合并者 AndreasKaratzas，指出的都是包裹时机边界问题，最终都被合并版本吸收：

• drafter 在 FULL target 下的 PIECEWISE 分发（正确性，已解决）：initialize_cudagraph_keys() 会把 FULL target 下的 drafter 映射为 PIECEWISE，若按中间版本（has_piecewise_cudagraphs）跳过 breakable 安装，drafter 将因 breakable 关闭编译而完全 eager 运行。合并版本对 drafter.model 无条件包裹 runtime_mode=PIECEWISE 的 breakable wrapper，规避该问题。
• FULL→FULL_AND_PIECEWISE 后期升级（设计，已解决）：显式 attention splitting_ops 时，FULL 可能在解析 UNIFORM_BATCH 后端后被升级为 FULL_AND_PIECEWISE，若安装已过则“选中 piecewise 却无捕获者”——即 #55095 想避免的回归。合并版本把 elif 改为独立 if 并保持 breakable 条件为 != NONE，保证 wrapper 在模式升级前就已就位。

• njhill 收到评论后回复 “Hopefully I have addressed them now”，随后 AndreaasKaratzas 给出 APPROVED 并确认 CI（含 AMD nightly）通过，仅剩一个已知 AMD 失败。

• FULL target 下 drafter 需要 breakable 包裹 (correctness): 合并版本对 drafter.model 无条件安装 runtime_mode=PIECEWISE 的 BreakableCUDAGraphWrapper，保证 drafter 的 PIECEWISE 分发始终有捕获者。

• FULL→FULL_AND_PIECEWISE 后期升级的包裹时机 (design): 合并版本将 elif 改为独立 if，且 breakable 条件保持 cudagraph_mode != NONE；wrapper 在模式升级前已安装，运行时按 runtime_mode=PIECEWISE 精确匹配，FULL 由外层标准 wrapper 捕获。

风险与影响

• 风险：
  1. 核心路径变更：load_model 是 MRV1 模型装配的必经之路，包裹顺序变化影响所有开启 VLLM_USE_BREAKABLE_CUDAGRAPH=1 的用户；FULL_AND_PIECEWISE 下新增双层嵌套使 dispatch 链路更长。
  2. 依赖 forward_context 的分发标注：breakable 是否捕获完全由 forward_context.cudagraph_runtime_mode 决定，任何调用方对模式标注不一致（如 NONE 与 PIECEWISE 混用）都会导致意外 eager 回退或捕获空转，形成隐性性能回归。
  3. 测试覆盖缺口：新增测试集中在 wrapper 单元层，缺少 FULL_AND_PIECEWISE 嵌套下的 e2e/集成测试（真实模型加载 + 捕获），模式升级路径 _check_and_update_cudagraph_mode 的行为只能靠代码阅读保证。
  4. unwrap() 的 hasattr 探测：若 runnable 恰好定义了同名 unwrap 属性且语义不同，递归会误触发；当前 wrapper 均为仓库内类，风险低但需留意第三方 runnable。
  5. 编译关闭的连带影响（评论提及 breakable 会关闭编译）：纯 FULL 模式下 drafter 被 breakable 化后其编译同样被关闭，需要确认该场景的性能预期。
• 影响：
  • 用户影响：仅 MRV1 + 显式开启 VLLM_USE_BREAKABLE_CUDAGRAPH=1 的部署受影响，行为从“整条路径 breakable”变为“FULL 用标准捕获、PIECEWISE 用 breakable”，与 MRV2 的 use_breakable_cg 行为对齐；未开启该开关的用户零影响。
  • 系统影响：FULL_AND_PIECEWISE 下出现两层 wrapper 嵌套，捕获条目由两个 wrapper 分别持有，内存与 graph pool 使用略增；但消除了 breakable 劫持 FULL 捕获带来的正确性风险。
  • 团队影响：确立 wrapper 嵌套契约（runtime_mode 严格门控 + 递归 unwrap），为后续在单一模型上混合使用 FULL 与 breakable 捕获提供可复用模式；MRV2 采用默认 match-any，两条 runner 路径的行为在此收敛。
  • 风险标记：核心路径变更, 嵌套 wrapper 契约, 缺少 e2e 测试覆盖

关联脉络

• PR #55095 前序 PR（标题未在材料中提供）: AndreaasKaratzas 在 review 评论中引用 #55095，说明其曾试图避免“选择 piecewise 却没有捕获者”的回归；本 PR 的包裹策略（breakable 先装 + 外层 FULL 再套）正是为杜绝该问题复发。