执行摘要

• 一句话：AR 投机解码多步融合进单个 CUDA graph，按 backend 能力门控
• 推荐动作：值得精读。这是 MRV2 投机解码路径的一次重要架构改进，核心看点有三个：一是以 backend 能力矩阵做特性门控并自动回退的安全设计；二是将“只在捕获期执行”的契约（capture-safe、持久存储）写进基类协议的文档化思路；三是作者对性能测量非常克制——既给出 A/B 数据，也承认 device-bound 场景无系统收益，并用 acceptance-normalized 指标排除采样率干扰。读者应重点关注 _generate_fused_drafts 的循环顺序与 update_draft_decode_metadata 的调用时机，以及 FA3/sparse SWA 两个适配示例。

功能与动机

PR body 明确说明动机：#41162 修复 stale attention metadata 的方式是每 draft 步重建 metadata 并重放独立 CUDA graph，虽然正确但重新引入了 per-step Python dispatch、metadata 构造和 CUDA graph launch 开销。作者还提供了历史调查证据：pre-#41162 版本即便加 metadata refresh 也无法复现修复，说明非法内存访问不完全由 stale draft metadata 解释，为直接恢复 fused 执行提供了依据。

实现拆解

第 1 步：投机器侧建立 fused 路径开关与 backend 能力门控

AutoRegressiveSpeculator（vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py）新增 use_fused_multi_step_decode 状态位，并在 set_attn() 中追加 _configure_fused_multi_step_decode()。该函数按三个优先级决策：单步投机直接关闭；advance_draft_positions=False 的位置静态草稿模型（如 Gemma4 MTP）天然安全、直接开启；其余情况扫描所有 attention group，只要有一个 backend 未声明 supports_draft_decode_metadata_update 就整体回退到逐 step 重建元数据的老路径，并打印 logger.info_once 提示。

第 2 步：attention backend 契约层新增刷新协议

AttentionMetadataBuilder 基类（vllm/v1/attention/backend.py）新增默认 False 的 supports_draft_decode_metadata_update 类属性，以及默认抛 NotImplementedError 的 update_draft_decode_metadata()；docstring 强调该方法只在 CUDA graph capture 期间调用，实现必须 emit capture-safe 操作并将重放态张量保存在持久存储。AttentionGroup（vllm/v1/worker/utils.py）透传该能力并路由刷新调用到组内第一个 metadata builder。

第 3 步：capture 与 propose 分流，实现 fused 执行循环

capture() 根据开关选择捕获 _generate_fused_drafts 或 _generate_draft；propose() 相应选择 _fused_multi_step_decode 或 _multi_step_decode。fused 路径先构建一次 slot mappings 和 step=1 的 draft attn metadata，FULL 模式下只 run_fullgraph 一次；非 FULL 模式走 _generate_fused_drafts，循环中每步执行 _generate_draft，非最后一步时重新计算 slot mappings 并对每个 attn group 调用 update_draft_decode_metadata()。附带修复：capture 前重置 idx_mapping（因为 #48892 后 padded 条目持久为 -1）。

第 4 步：按 backend 逐个适配并声明能力

FlashAttention（vllm/v1/attention/backends/flash_attn.py）将 build() 内的 schedule 闭包重构为 _get_scheduler_metadata() / _store_scheduler_metadata() 两个方法（metadata 写入持久 buffer 的语义保留），新增 update_draft_decode_metadata() 重算 FA3 AOT scheduler metadata；能力位设为 dcp_world_size == 1，DCP 场景因 host-side 决策可能在 replay 间改变控制流而显式排除。DeepSeek sparse SWA（vllm/v1/attention/backends/mla/sparse_swa.py）的 update_draft_decode_metadata() 用 Triton kernel 原地重算 SWA indices/lens、重建 tile scheduler 并清空 FlashInfer sparse index cache。Triton Attention 与 Triton MLA 声明支持（其步进相关字段已引用持久输入 buffer）。ROCm 的 DeepseekV4ROCMAiterSparseSWAMetadataBuilder（vllm/models/deepseek_v4/amd/rocm.py）显式关闭，待 ROCm 特定 ragged SWA 适配。

第 5 步：测试与文档配套

tests/v1/worker/test_gpu_autoregressive_speculator.py 新增参数化测试，断言 _multi_step_decode 与 _fused_multi_step_decode 在 CUDAGraphMode.NONE/FULL 下 _generate_draft 的 eager 调用次数与 run_fullgraph 次数（fused+FULL 为 0 eager / 1 replay，split+FULL 为 0 / 3）；另新增 FA3 scheduler metadata 刷新的正向与跳过单测。docs/design/model_runner_v2.md 新增 Fused Multi-Step Draft Decoding 章节，说明新契约的启用条件与实现约束。

关键文件：

• vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py（模块 投机解码；类别 source；类型 core-logic；符号 set_attn, _configure_fused_multi_step_decode, _fused_multi_step_decode, _generate_fused_drafts）: 核心变更文件：新增 fused 多步解码的开关决策、capture/propose 分流与完整的 fused 执行循环，承载了本 PR 的主要执行逻辑。
• vllm/v1/attention/backends/flash_attn.py（模块 注意力后端；类别 source；类型 core-logic；符号 _get_scheduler_metadata, _store_scheduler_metadata, schedule, update_draft_decode_metadata）: FA3 是首个需要真正刷新派生元数据的 backend：将 schedule 闭包重构为可复用方法，并实现 update_draft_decode_metadata 重算 AOT scheduler metadata。
• tests/v1/worker/test_gpu_autoregressive_speculator.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_multi_step_decode_replays_captured_graph_as_expected, test_update_draft_decode_metadata_updates_fa3_scheduler_metadata, fake_get_scheduler_metadata, test_update_draft_decode_metadata_skips_without_scheduler_metadata）: 新增参数化测试覆盖 fused/split 两种路径在 FULL 与 NONE 模式下的 graph replay 次数，以及 FA3 元数据刷新的正向与跳过分支。
• vllm/v1/attention/backends/mla/sparse_swa.py（模块 注意力后端；类别 source；类型 core-logic；符号 update_draft_decode_metadata）: DeepSeek V4 sparse SWA 的刷新实现最复杂：重算 SWA indices/lens、重建 tile scheduler，并清空 FlashInfer 稀疏索引缓存。
• vllm/v1/worker/utils.py（模块 工具层；类别 source；类型 core-logic；符号 supports_draft_decode_metadata_update, update_draft_decode_metadata）: AttentionGroup 层透传能力声明与刷新调用，是 speculator 与具体 metadata builder 之间的路由枢纽。
• vllm/v1/attention/backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 update_draft_decode_metadata）: 基类定义新契约：supports_draft_decode_metadata_update 默认关闭与 update_draft_decode_metadata 默认抛错，强制所有 backend 显式适配。
• vllm/v1/attention/backends/mla/triton_mla.py（模块 注意力后端；类别 source；类型 core-logic；符号 update_draft_decode_metadata）: Triton MLA 声明支持 fused 路径（非 DCP 时），刷新方法为空实现。
• vllm/v1/attention/backends/triton_attn.py（模块 注意力后端；类别 source；类型 core-logic；符号 update_draft_decode_metadata）: Triton Attention 可直接声明支持，因为步进相关字段已引用持久输入 buffer，无需额外刷新。
• vllm/models/deepseek_v4/amd/rocm.py（模块 模型适配；类别 source；类型 data-contract）: ROCm sparse SWA 因 ragged SWA indices/indptrs 尚未适配而显式禁用 fused 路径，避免回归。
• docs/design/model_runner_v2.md（模块 设计文档；类别 docs；类型 documentation）: 补充 fused multi-step draft decoding 的设计文档，说明契约、启用条件与实现约束。

关键符号：_configure_fused_multi_step_decode, _fused_multi_step_decode, _generate_fused_drafts, update_draft_decode_metadata, _get_scheduler_metadata, _store_scheduler_metadata

关键源码片段

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

核心变更文件：新增 fused 多步解码的开关决策、capture/propose 分流与完整的 fused 执行循环，承载了本 PR 的主要执行逻辑。

# vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
# 依据各 attention backend 的能力，决定是否启用 fused 多步草稿解码。
def _configure_fused_multi_step_decode(self) -> None:
    # 单步投机没有“多步”可融合，直接关闭 fused 路径。
    if self.num_speculative_steps == 1:
        self.use_fused_multi_step_decode = False
        return
​
    # 位置不推进的草稿模型（如 Gemma4 MTP）不存在位置相关元数据
    # 的失效问题，因此天然可以安全融合。
    if not self.advance_draft_positions:
        self.use_fused_multi_step_decode = True
        return
​
    # 只要有一个 attention group 不支持 draft decode 元数据原地刷新，
    # 就整体回退到逐 step 重建元数据的老路径，保证正确性优先。
    unsupported_backends = sorted({
        attn_group.backend.get_name()
        for attn_groups in self.attn_groups
        for attn_group in attn_groups
        if not attn_group.supports_draft_decode_metadata_update
    })
    self.use_fused_multi_step_decode = not unsupported_backends
    if unsupported_backends:
        logger.info_once(
            "Fused multi-step draft decode is not supported by attention "
            "backend(s) %s; falling back to rebuilding attention metadata "
            "between draft steps.",
            ", ".join(unsupported_backends),
        )
​
​
# 单个 CUDA graph 内捕获全部 post-prefill 草稿步骤。
# capture 时本方法被录制；replay 时整段作为图执行，Python 不再介入。
def _generate_fused_drafts(
    self,
    num_reqs: int,
    num_tokens_padded: int,
    attn_metadata: dict[str, Any] | None,
    slot_mappings: dict[str, torch.Tensor] | None,
    num_tokens_across_dp: torch.Tensor | None,
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
) -> None:
    idx_mapping = self.idx_mapping[:num_reqs]
    positions = self.input_buffers.positions[:num_reqs]
    query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
    attn_groups = (
        [group for groups in self.attn_groups for group in groups]
        if attn_metadata is not None
        else []
    )
​
    for step in range(1, self.num_speculative_steps):
        self.current_draft_step.fill_(step)
        # 每步执行草稿模型 forward + sampling（capture 时记录）。
        self._generate_draft(
            num_reqs, num_tokens_padded, attn_metadata, slot_mappings,
            num_tokens_across_dp, cudagraph_runtime_mode,
        )
        # 非最后一步时推进输入位置并原地刷新 backend 派生元数据；
        # update_draft_decode_metadata() 仅影响图捕获期，replay 时
        # 其 GPU 操作已固化在图中。
        if (step < self.num_speculative_steps - 1
                and attn_metadata is not None
                and self.advance_draft_positions):
            self.block_tables.compute_slot_mappings(
                idx_mapping, query_start_loc, positions, num_tokens_padded,
            )
            for attn_group in attn_groups:
                attn_group.update_draft_decode_metadata(attn_metadata)

vllm/v1/attention/backends/flash_attn.py

FA3 是首个需要真正刷新派生元数据的 backend：将 schedule 闭包重构为可复用方法，并实现 update_draft_decode_metadata 重算 AOT scheduler metadata。

# vllm/v1/attention/backends/flash_attn.py
# FA3 的 scheduler metadata 是 position-dependent 的派生状态，
# 在 fused 多步草稿解码中必须随 draft 步骤原地刷新。
def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None:
    # 非 FA3 或没有 scheduler metadata 时无需处理。
    if metadata.scheduler_metadata is None:
        return
​
    num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0]
​
    # fused 路径仅在无 DCP 时启用（见 __init__ 中按 dcp_world_size 门控）；
    # draft decode 恒以 common_prefix_len=0 构建，因此不存在 cascade。
    assert self.dcp_world_size == 1
    assert not metadata.use_cascade
​
    # 重新计算 AOT scheduler metadata，并写入持久 buffer，
    # 保证 CUDA graph replay 期间地址稳定。
    scheduler_metadata = self._get_scheduler_metadata(
        aot_schedule=True,
        batch_size=num_reqs,
        cu_query_lens=metadata.query_start_loc,
        max_query_len=metadata.max_query_len,
        seqlens=metadata.seq_lens,
        max_seq_len=metadata.max_seq_len,
        causal=metadata.causal,
        max_num_splits=metadata.max_num_splits,
    )
    metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
​
​
# 与 build() 共用：写入持久 buffer 并清零多余槽位，避免无效
# scheduler metadata 被线程块使用而覆盖输出 buffer。
def _store_scheduler_metadata(
    self, scheduler_metadata: torch.Tensor | None
) -> torch.Tensor | None:
    if self.use_full_cuda_graph and scheduler_metadata is not None:
        n = scheduler_metadata.shape[0]
        assert self.scheduler_metadata is not None
        self.scheduler_metadata[:n] = scheduler_metadata
        self.scheduler_metadata[n:] = 0
        return self.scheduler_metadata[:n]
    return scheduler_metadata

vllm/v1/attention/backends/mla/sparse_swa.py

DeepSeek V4 sparse SWA 的刷新实现最复杂：重算 SWA indices/lens、重建 tile scheduler，并清空 FlashInfer 稀疏索引缓存。

# vllm/v1/attention/backends/mla/sparse_swa.py
# DeepSeek V4 sparse SWA 的草稿元数据刷新：
# 重算 SWA lengths/indices、重建 tile scheduler、并使 FlashInfer 稀疏索引缓存失效。
def update_draft_decode_metadata(
    self,
    metadata: DeepseekSparseSWAMetadata,
) -> None:
    if metadata.num_decode_tokens == 0:
        return
    assert metadata.query_start_loc is not None
    assert metadata.seq_lens is not None
    assert metadata.token_to_req_indices is not None
    assert metadata.is_valid_token is not None
    assert metadata.decode_swa_indices is not None
    assert metadata.decode_swa_lens is not None
​
    # 原地重算 SWA 索引和长度；该 Triton kernel 调用在 CUDA graph capture
    # 期间被记录，replay 时随图执行，不会回到 Python。
    _compute_swa_indices_and_lens_kernel[(metadata.num_decode_tokens,)](
        metadata.decode_swa_indices,
        metadata.decode_swa_indices.stride(0),
        metadata.decode_swa_lens,
        metadata.decode_swa_indices.shape[-1],
        metadata.query_start_loc,
        metadata.seq_lens,
        metadata.token_to_req_indices,
        metadata.is_valid_token,
        metadata.block_table,
        metadata.block_table.stride(0),
        self.block_size,
        token_offset=0,
        TRITON_BLOCK_SIZE=1024,
    )
    # 每个 draft 步骤的 tile scheduler 需要重新规划。
    tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)
    metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
    metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]
    metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]
    # review 中发现的遗漏：该缓存是在 _forward 中构建的，不清空会
    # 复用第一步的 stale 稀疏索引，导致后续草稿步错误。
    metadata.flashinfer_sparse_index_cache.clear()

评论区精华

Review 核心交锋集中在正确性与契约设计：TheEpicDolphin 要求以 backend 能力门控 fused 路径，指出 FlashAttention 的 scheduler metadata 是 position-dependent 的，在 draft 循环中会变 stale；作者以 supports_draft_decode_metadata_update 增量适配回应。讨论中修复了一处真实正确性缺陷——flashinfer_sparse_index_cache 未在 draft 步骤间失效，作者承认该状态在 _forward 而非 builder 中构建导致遗漏。另一个关键点是 cudagraph-safe：FA3 cascade 的 prefix_scheduler_metadata 新对象赋值被判定不安全后移除。设计层面，reviewer 建议移除无意义的 enable_fused_decode_graph 配置 flag（作者照做），并提议对齐 update_seq_lens 命名（作者以刷新语义更宽为由保留现名）。

• 用 backend 能力门控 fused 路径，避免 stale 派生元数据 (design): 通过 supports_draft_decode_metadata_update 逐 backend 声明，任一不支持则整体回退到逐 step 重建路径。
• flashinfer_sparse_index_cache 未失效导致 stale 稀疏索引 (correctness): 在 update_draft_decode_metadata() 末尾加入 metadata.flashinfer_sparse_index_cache.clear()。
• 配置 flag enable_fused_decode_graph 是否必要 (design): 作者移除了该公开配置选项，fused 路径完全由 backend 能力自动选择。
• FA3 cascade 的 prefix_scheduler_metadata 新对象赋值非 cudagraph-safe (correctness): 移除 cascade 分支，并断言 not metadata.use_cascade。
• update_draft_decode_metadata 必须保证 CUDA graph capture 安全 (documentation): 基类 docstring 补充说明：实现必须使用 capture-safe 操作，重放态张量必须持久存储。
• 接口命名：update_seq_lens vs update_draft_decode_metadata (style): 保留 update_draft_decode_metadata 命名。

风险与影响

• 风险：
  1. 捕获期正确性风险：fused 路径把整个草稿循环录制进单一 CUDA graph，任何由 backend 在 forward 中派生、未纳入刷新协议的元数据都会在 replay 时 stale。本 PR review 已发现并修复 flashinfer_sparse_index_cache 一处，但 audit 面还可能存在其他类似状态（如 ROCm ragged SWA 尚未适配，因此被显式禁用）。
  2. DCP 排除：FA3 能力位硬编码为 dcp_world_size == 1，上下文并行下 fused 路径整体关闭，未来启用需设计完整的 replay-safe 刷新模型。
  3. 内存与捕获开销：整循环一次捕获会让捕获期 graph buffer 更大，capture 时间与显存占用上升；device-bound 场景下收益不确定。
  4. 行为窗口：_fused_multi_step_decode 在非 FULL 模式也走 fused 循环（reviewer 明确要求），意味着 eager 路径同样要满足捕获安全的约束，测试仅覆盖 FULL 与 NONE 两个模式端点。
     - 影响：对用户：MTP 投机在 host-bound 负载（小 TP、低并发、长草稿链）下每 scheduler step 的 graph launch 从 4 次降到 3 次，draft propose CPU 开销下降约 50%，profiler step CPU span 下降 29%；device-bound 在线服务（H100 8x80G 高并发）ITL 与 acceptance 均无系统性变化。对系统：MRV2 投机解码执行模型从“每步一个 graph”变为“整循环一个 graph”，减少 Python 侧 dispatch 与 metadata 构造。对团队：建立 backend 能力矩阵契约，所有 attention backend（含未来新增）都需声明 supports_draft_decode_metadata_update 并保证 update_draft_decode_metadata() 的 capture-safety，增加了 backend 适配成本。
     - 风险标记：核心路径变更, cudagraph 捕获期逻辑, backend 协议扩展, DCP 场景被排除, 性能收益依赖 host-bound 场景

关联脉络

• PR #41162 Fix stale draft attention metadata: PR body 明确引用：该 PR 修复 stale attention metadata 但引入 per-step graph 重建开销，本 PR 在其基础上恢复 fused multi-step graph 执行，是直接前置依赖。
• PR #48892 Persist padded idx_mapping entries: commit 0fe9ff4 提到本 PR 需在 capture 前重置 idx_mapping，因为 #48892 之后 padded idx_mapping 持久为 -1。
• PR #50493 [Kimi-K3] support DCP partial prefix cache hit: 同属 MRV2 路径的 DCP 演进；本 PR 明确将 DCP 场景排除在 fused 路径外（supports_draft_decode_metadata_update = dcp_world_size == 1），二者在 DCP 元数据刷新模型上互补。