执行摘要

• 一句话：CUDA 图捕获始终构建注意力元数据
• 推荐动作：此 PR 是一个值得精读的修复，展示了 CUDA 图捕获中注意力元数据管理的设计决策。注释中清晰解释了 FULL 和 PIECEWISE 模式下元数据构建的权衡，对理解 vLLM 的 CUDA 图架构有参考价值。

功能与动机

在 CUDA 图捕获中，注意力元数据不仅被标准注意力核需要，也被专门的注意力类操作（如 Inkling 的 sconv、DSV4 压缩器）所依赖。这些操作维护自身状态并需要 block table 等特殊元数据。此前 PIECEWISE 模式构建 CUDA 图时跳过了元数据构建，导致这些操作在捕获时无法正确执行。此变更通过始终构建元数据来修复此问题。

实现拆解

1. 修改 prepare_inputs_to_capture 函数签名：将 skip_attn: bool 参数替换为 full_cudagraph: bool，使调用者能够显式声明是否正在为 FULL CUDA 图构建元数据。

2. 统一构建注意力元数据：移除 skip_attn 分支，始终调用 model_state.prepare_attn() 构建元数据。当 full_cudagraph=True 时设置 for_capture=True，为 FULL 图生成可捕获的元数据；当 full_cudagraph=False（PIECEWISE 模式）时设置 for_capture=False，确保注意力类操作正常执行但标准注意力核不会生成不兼容的元数据。

3. 更新 cudagraph_utils.py 和规范解码模块中的调用处：在 CUDAGraphModelRunner.capture() 和 AutoregressiveCUDAGraphRunner.capture() 中，将 skip_attn 参数替换为 full_cudagraph，同时移除与 PIECEWISE 模式相关的 assert 断言。

4. 包含详细的注释说明：在 prepare_inputs_to_capture 中添加了大量注释，解释了 FULL 和 PIECEWISE 模式下构建元数据的设计决策和注意事项。

关键文件：

• vllm/v1/worker/gpu/cudagraph_utils.py（模块 调度器；类别 source；类型 core-logic；符号 prepare_inputs_to_capture, CUDAGraphModelRunner.capture）: 核心变更文件，修改了 CUDA 图捕获过程中注意力元数据的构建逻辑。将 skip_attn 参数替换为 full_cudagraph，并始终调用 prepare_attn 构建元数据。添加了详细注释说明设计决策。
• vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py（模块 调度器；类别 source；类型 core-logic；符号 AutoregressiveCUDAGraphRunner.capture）: 兄弟文件，在规范解码的 CUDA 图捕获路径中进行了类似的参数替换，保持了统一的行为。

关键符号：prepare_inputs_to_capture, CUDAGraphModelRunner.capture, AutoregressiveCUDAGraphRunner.capture

关键源码片段

vllm/v1/worker/gpu/cudagraph_utils.py

核心变更文件，修改了 CUDA 图捕获过程中注意力元数据的构建逻辑。将 skip_attn 参数替换为 full_cudagraph，并始终调用 prepare_attn 构建元数据。添加了详细注释说明设计决策。

# vllm/v1/worker/gpu/cudagraph_utils.py
​
def prepare_inputs_to_capture(
    num_reqs: int,
    num_tokens: int,
    model_state: ModelState,
    input_buffers: InputBuffers,
    block_tables: BlockTables,
    attn_groups: list[list[AttentionGroup]],
    kv_cache_config: KVCacheConfig,
    full_cudagraph: bool, # 替代原有的 skip_attn: bool
) -> AttentionState:
    # ... 创建 dummy input_batch 和 block_tables ...
    input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
    input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
    # ... 构建 slot mappings ...
​
    # NOTE(woosuk): Attention metadata 不仅被标准 attention 核需要 ,
    # 也被专门的 attention-like 操作 ( 如 Inkling 的 sconv, DSV4 压缩器 )
    # 所依赖 , 这些操作维护自身状态并需要 block table 等特殊 metadata.
    # 因此 , 我们始终构建 attention metadata, 而不是跳过 .
​
    # - 对于 FULL CUDA graph: 设置 for_capture=True, 使得 attention 和
    # attention-like 操作都能产生可捕获的 metadata.
    # - 对于 PIECEWISE CUDA graph: 设置 for_capture=False, 因为 :
    # * attention-like 操作可能不作为断点 , 所以需要 metadata 来执行和捕获 .
    # * 标准 attention 作为断点时会 eager 执行 , for_capture=False 是必需的 ,
    # 因为某些后端 ( 如 linear attention) 无法为 prefill 生成可捕获的 metadata.
    # 这里假设 attention-like 操作即使 for_capture=False 也能产生可捕获 metadata.
    attn_metadata = model_state.prepare_attn(
        input_batch,
        CUDAGraphMode.NONE,
        input_block_tables,
        slot_mappings,
        attn_groups,
        kv_cache_config,
        for_capture=full_cudagraph,
    )
​
    # ... 返回 attn_metadata 和 slot_mappings ...
    return attn_metadata, slot_mappings

vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py

兄弟文件，在规范解码的 CUDA 图捕获路径中进行了类似的参数替换，保持了统一的行为。

# vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
​
def create_forward_fn(desc, warmup):
    # ... 创建 num_tokens, num_reqs ...
    attn_metadata, slot_mappings = prepare_inputs_to_capture(
        num_reqs,
        num_tokens,
        model_state,
        input_buffers,
        block_tables,
        attn_groups,
        kv_cache_config,
        # 将原来基于 skip_attn 的条件改为显式的 full_cudagraph 参数
        full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
    )
​
    return lambda cg_mode: forward_fn(
        num_reqs,
        num_tokens,
        attn_metadata,
        slot_mappings,
        num_tokens_across_dp,
        cg_mode,
    )

评论区精华

该 PR 是 @WoosukKwon 此前 PR #49364 的重新提交，之前因暴露了 MRV2 中的一个 bug 导致 CI 失败而被回滚。相关 bug 已在 PR #49751 中修复，因此本 PR 现在可以顺利合并。审核由 WoosukKwon 本人批准，且没有额外的讨论线程。

• 暂无高价值评论线程

风险与影响

• 风险：主要风险在于始终构建注意力元数据可能引入不必要的开销。但代码中通过 for_capture 参数确保了 FULL 和 PIECEWISE 模式下元数据的正确性，且此前 PR #49364 已被验证（仅因另外一个独立 bug 被回滚）。如果某个注意力后端在 for_capture=False 时生成了不兼容的元数据，可能导致 CUDA 图捕获失败。但现有测试和注释中提到该假设在实践中有效。
• 影响：影响范围限于 MRV2 的 CUDA 图捕获流程。对用户无直接可见影响，但修复了 PIECEWISE CUDA 图中注意力类操作（如 Inkling sconv、DSV4 压缩器）无法工作的问题。变更代码量小，涉及两个文件。
• 风险标记：核心路径变更, 依赖修复, 假设存在脆弱性(attention-like ops at capture)

关联脉络

• PR #49364 [MRV2] Always build attn metadata at capture time: 本 PR 是 PR #49364 的重新提交，原始 PR 曾因暴露的 bug 被回滚。
• PR #49751 [MRV2] Fix bug exposed by always building attn metadata: 修复了 PR #49364 暴露的 MRV2 bug，使得本 PR 能够安全合并。