执行摘要

• 一句话：为 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 的独立性

# 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 图支持。