# PR #48822 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Model] Add PW CUDA graph support for Inkling [2/N]
- 合并时间：2026-07-17 01:28
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/48822

---

# 执行摘要

- 一句话：为 Inkling 模型添加分段 CUDA 图支持
- 推荐动作：建议相关开发者详细阅读 `cudagraph_utils.py` 中 capture 逻辑的变化，理解 breakable 图与完整图在 metadata 管理上的差异。测试函数 `test_piecewise_capture_builds_fresh_metadata_for_both_passes` 展示了关键行为，值得学习。整体设计采用条件分支模式，可复用于未来其他需要 breakable CUDA 图的模型。由于改动涉及 CUDA 图管理核心，合并后建议密切监控相关模型（尤其是 Inkling）的性能和正确性。

# 功能与动机

这是从 #48768 中切出的第二块独立可合并部分。Inkling 模型需要在 Model Runner V2 上支持 breakable CUDA 图以提升解码性能，同时保持注意力元数据的正确性。PR 动机在于为 Inkling 提供 PIECEWISE CUDA 图支持，利用 breakable 图避免完整图的请求填充限制。

# 实现拆解

1. **调整 Inkling 注意力后端 CUDA 图支持等级**：在 `vllm/models/inkling/nvidia/sconv_swa_attn.py` 中将 `_cudagraph_support` 从 `UNIFORM_SINGLE_TOKEN_DECODE` 改为 `UNIFORM_BATCH`，允许 PIECEWISE 模式。

2. **修改 CudaGraphManager 的 capture 逻辑 **（`vllm/v1/worker/gpu/cudagraph_utils.py`）：
 - 将 `_init_candidates` 调用后移至 `use_breakable_cg` 设置之后，确保候选依赖正确状态。
 - 在 `_init_candidates` 中，当 `mix_mode == CUDAGraphMode.PIECEWISE` 且 `use_breakable_cg` 时，将 `num_reqs` 设为 `min(num_tokens, self.max_num_reqs)`（非 breakable 时为 `None`）。
 - 在 `capture` 方法中，根据 `use_breakable_cg` 决定 PIECEWISE 描述符的 `forward_fn` 调用：若使用 breakable，则通过 `create_forward_fn` 获取 fresh metadata 再调用；否则直接调用已有的 `forward_fn`（跳过 attention）。
 - 同步调整 `skip_attn` 参数和 `attn_metadata` 断言。

3. **调整 DefaultModelState.prepare_attn 填充逻辑 **（`vllm/v1/worker/gpu/model_states/default.py`）：当 `cudagraph_mode == CUDAGraphMode.PIECEWISE` 且 `is_breakable_cudagraph_enabled()` 时，使用填充后的 `num_reqs` 和 `num_tokens`，确保 attention metadata 尺寸在 capture 和 replay 间一致。

4. **适配 spec_decode 的 CUDA 图工具 **（`vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py`）：将 `skip_attn` 条件同样调整为只在非 breakable 时跳过。

5. **新增与修改测试**：
 - `tests/v1/cudagraph/test_breakable_cudagraph.py`：新增 `test_piecewise_capture_builds_fresh_metadata_for_both_passes`，验证 capture 过程中 warmup 和 capture 使用独立的 metadata 对象。
 - `tests/models/inkling/test_contract_validation.py`：将原有 `test_inkling_supports_full_decode_only_cudagraphs` 改为 `test_inkling_supports_piecewise_cudagraphs`，断言支持 PIECEWISE 模式且 resolve 后返回 `CUDAGraphMode.PIECEWISE`。

关键文件：
- `vllm/v1/worker/gpu/cudagraph_utils.py`（模块 图管理；类别 source；类型 core-logic；符号 CudaGraphManager.capture, CudaGraphManager._init_candidates, CudaGraphManager.__init__）: 核心修改：调整 capture 和 _init_candidates 以支持 breakable PIECEWISE 模式
- `tests/v1/cudagraph/test_breakable_cudagraph.py`（模块 图测试；类别 test；类型 test-coverage；符号 test_piecewise_capture_builds_fresh_metadata_for_both_passes, create_forward_fn, forward_fn）: 新增测试验证 PIECEWISE capture 中 metadata 的独立性
- `vllm/v1/worker/gpu/model_states/default.py`（模块 状态层；类别 source；类型 data-contract；符号 DefaultModelState.prepare_attn）: 调整 prepare_attn 在 PIECEWISE 且 breakable 启用时使用填充尺寸
- `tests/models/inkling/test_contract_validation.py`（模块 契约测试；类别 test；类型 test-coverage；符号 test_inkling_supports_piecewise_cudagraphs, test_inkling_supports_full_decode_only_cudagraphs）: 修改测试以验证 Inkling 注意力模块支持 PIECEWISE 模式
- `vllm/models/inkling/nvidia/sconv_swa_attn.py`（模块 注意力层；类别 source；类型 data-contract；符号 InklingSconvMetadataBuilder._cudagraph_support）: 修改 Inkling 注意力后端的 CUDA 图支持等级以允许 PIECEWISE

关键符号：CudaGraphManager.capture, CudaGraphManager._init_candidates, DefaultModelState.prepare_attn, InklingSconvMetadataBuilder.get_cudagraph_support, test_piecewise_capture_builds_fresh_metadata_for_both_passes, test_inkling_supports_piecewise_cudagraphs

## 关键源码片段

### `tests/v1/cudagraph/test_breakable_cudagraph.py`

新增测试验证 PIECEWISE capture 中 metadata 的独立性

```python
# tests/v1/cudagraph/test_breakable_cudagraph.py ( 新增测试函数 )
def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
    """验证 PIECEWISE capture 过程中 warmup 和 capture 使用不同的 metadata 对象。"""
    from vllm.config import CUDAGraphMode
    from vllm.v1.worker.gpu.cudagraph_utils import (
        BatchExecutionDescriptor,
        CudaGraphManager,
    )

    # 手动构造 CudaGraphManager 实例（跳过 __init__ 以避免复杂依赖）
    manager = CudaGraphManager.__new__(CudaGraphManager)
    desc = BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 8, None)
    manager.device = torch.device("cpu")
    manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]}
    manager._graphs_captured = False
    manager.use_breakable_cg = True

    create_calls = []  # 记录 create_forward_fn 的调用参数
    forward_calls = []  # 记录 forward_fn 的调用参数

    def create_forward_fn(desc_arg, warmup):
        """模拟模型的状态初始化：每次调用创建独立的 metadata。"""
        assert desc_arg == desc
        metadata = {"layer": object()}
        create_calls.append((warmup, metadata))

        def forward_fn(cg_mode):
            """模拟前向传播，使用闭包捕获的 metadata。"""
            assert metadata  # 确保 metadata 存在
            forward_calls.append((warmup, cg_mode, metadata))

        return forward_fn

    # mock graph_capture 和 is_global_first_rank 以避免 CUDA 依赖
    with (
        patch(
            "vllm.v1.worker.gpu.cudagraph_utils.graph_capture",
            return_value=nullcontext(),
        ),
        patch(
            "vllm.v1.worker.gpu.cudagraph_utils.is_global_first_rank",
            return_value=False,
        ),
    ):
        manager.capture(create_forward_fn)

    # 验证：先调用 warmup=True，后调用 warmup=False
    assert [warmup for warmup, _ in create_calls] == [True, False]
    # 验证 forward_fn 的 cg_mode 顺序：先 NONE（warmup），后 PIECEWISE（capture）
    assert [mode for _, mode, _ in forward_calls] == [
        CUDAGraphMode.NONE,
        CUDAGraphMode.PIECEWISE,
    ]
    # 验证两次产生的 metadata 对象相互独立
    assert create_calls[0][1] is not create_calls[1][1]

```

# 评论区精华

PR 在合并前未收到人工 review 评论，仅有机器人自动提醒。所有设计决策在 PR body 中明确说明，包括与 #48768 的切分关系。核心权衡在于 breakable 模式下 metadata 的独立创建与 attention 是否跳过。团队已通过模型评估验证正确性。

- 无人工 review (other): 无

# 风险与影响

- 风险：主要风险在于 breakable CUDA 图路径对非 Inkling 模型的影响。由于条件判断均通过 `use_breakable_cg` 和 `is_breakable_cudagraph_enabled()` 保护，默认为 False，不影响现有流程。但对于启用 breakable 的模型，需要确保 `AttentionCGSupport.UNIFORM_BATCH` 与模型实际能力匹配。Inkling 注意力模块的 metadata 构建在填充和非填充路径下均已验证。此外，`CudaGraphManager` 的 `_init_candidates` 调用顺序调整可能影响候选列表构建，但经测试未发现问题。新增测试覆盖了 metadata 独立性的核心场景，降低了回归风险。
- 影响：对 Inkling 模型用户，该 PR 允许启用 PIECEWISE CUDA 图，从而减少请求填充开销，提升解码吞吐。对系统，`cudagraph_utils.py` 的修正是通用的，但通过开关隔离，不影响其他模型。对团队，该 PR 是 Inkling 支持的第二块，后续还需 MTP 和 speculative decoding 支持。当前版本不包含 LoRA 和 Inkling MTP 的 CUDA 图支持，需注意适用范围。
- 风险标记：核心路径变更 , 模型特定优化 , 需启用 breakable 标志

# 关联脉络

- PR #48768 Inkling CUDA graph support (original): 此 PR 是从 #48768 切出的第二块，包含 breakable CUDA 图部分。
- PR #48799 [Model] Add Inkling model support [1/N]: 第一块 Inkling 模型支持，此 PR 在其基础上增加 CUDA 图支持。