执行摘要

• 一句话：DCP 下为 DSpark 草稿启用 FlashInfer MLA 支持
• 推荐动作：值得精读。重点关注三个设计决策：共享 draft metadata 构建器对 DCP 本地长度的统一处理、AttentionSpec.merge 从 any() 到组内一致的语义收紧、以及 FlashInfer 非因果多 token 展开时对全局因果边界的 repeat_interleave。建议在合并后尽快补跑 Kimi-K3 DSpark+DCP 端到端评估，并排查仓库内是否存在因果/非因果层共享 KV cache group 的存量配置。

功能与动机

DSpark 需要在 DCP 下使用 FlashInfer MLA，但现有实现存在两个互补缺口：其一，draft 注意力必须从自己的当前全局序列长度重算 rank-local 长度，因为 draft 步骤在采样和拒绝后会前进或回退 input_buffers.seq_lens，复用 target-local 长度不够；其二，FlashInfer 非因果路径把每个查询块展开成独立 decode 行，DCP 下局部 block table 与序列长度已展开，但全局因果边界没有，导致 native CP 内核批量几何不匹配。此外，cache 组合并使用 any(non_causal_multi_token_decode) 可能把因果 target 组提升为非因果行为，需要改为组内一致。

实现拆解

1. 共享 draft 元数据构建器重算 DCP 本地长度：在 vllm/v1/worker/gpu/spec_decode/speculator.py 的 _build_draft_attn_metadata 中新增分支，当 dcp_local_seq_lens 未传入且 block_tables.cp_size > 1 时，基于 input_buffers.seq_lens 调用 cp_utils.prepare_dcp_local_seq_lens 重算 rank-local 长度并复用给 build_attn_metadata。原因是 draft 步骤推进/回退 seq_lens 后，target 侧长度已过期。该逻辑位于共享基类，对 EAGLE、MTP、DFlash 均生效。
2. 移除 DFlash 重复准备逻辑：vllm/v1/worker/gpu/spec_decode/dflash/speculator.py 中删除与基类重复的 prepare_dcp_local_seq_lens 导入和调用（净删 11 行），改为直接委托父类，保持单一职责并消除行为分叉。
3. 收紧 KV cache group 合并语义：vllm/v1/kv_cache_interface.py 的 AttentionSpec.merge 将 non_causal_multi_token_decode 从 any() 改为 set 去重后断言唯一，确保因果 target 组与非因果 draft 组不会被错误合并成一个 group。
4. FlashInfer MLA backend 能力声明与几何展开：vllm/v1/attention/backends/mla/flashinfer_mla.py 声明 supports_non_causal_multi_token_dcp = True，并在 forward_mqa 非因果多 token 路径对 causal_seqlens_kv_global 执行 repeat_interleave(query_len)，使全局因果边界与展开后的 decode 行几何一致。
5. 测试配套：tests/v1/spec_decode/test_eagle_draft_attn_metadata.py 新增 test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens 回归；tests/v1/attention/test_flashinfer_mla_dcp.py 将 test_flashinfer_mla_forward_uses_native_dcp_api 参数化为 causal/noncausal 两档；test_mla_noncausal.py 与 test_mla_backends.py 同步更新组语义与 DSpark DCP 支持断言。共 12 个 B300 单元测试。

关键文件：

• vllm/v1/worker/gpu/spec_decode/speculator.py（模块 投机解码；类别 source；类型 dependency-wiring；符号 _build_draft_attn_metadata, prepare_dcp_local_seq_lens）: 核心变更点：在共享 draft 元数据构建器中根据 draft 模型当前全局序列长度重算 rank-local DCP 长度，所有 speculator（EAGLE/MTP/DFlash）统一受益。
• vllm/v1/worker/gpu/spec_decode/dflash/speculator.py（模块 投机解码；类别 source；类型 dependency-wiring；符号 _build_draft_attn_metadata）: 移除 DFlash 重复的 DCP 本地长度准备逻辑，委托共享基类处理，消除行为分叉。
• tests/v1/spec_decode/test_eagle_draft_attn_metadata.py（模块 草稿元数据；类别 test；类型 test-coverage；符号 test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens, fake_prepare）: 新增 DCP 本地序列长度重算回归测试，通过 fake_prepare mock 验证 prepare_dcp_local_seq_lens 被正确调用且结果传给 build_attn_metadata。
• tests/v1/attention/test_flashinfer_mla_dcp.py（模块 MLA 测试；类别 test；类型 test-coverage；符号 test_flashinfer_mla_forward_uses_native_dcp_api）: 将 DCP native API 测试参数化为 causal/noncausal，验证非因果块展开后 block_table、seq_lens 与全局因果边界均按 query_len 重复。
• vllm/v1/attention/backends/mla/flashinfer_mla.py（模块 注意力后端；类别 source；类型 core-logic；符号 FlashInferMLAImpl.forward_mqa, FlashInferMLAMetadataBuilder.supports_non_causal_multi_token_dcp）: 核心 backend 变更：声明非因果多 token DCP 支持，并在 forward_mqa 中对全局因果边界做几何展开。
• vllm/v1/kv_cache_interface.py（模块 缓存接口；类别 source；类型 core-logic；符号 AttentionSpec.merge）: 收紧 AttentionSpec.merge 对 non_causal_multi_token_decode 的一致性要求，避免因果 target 组与非因果 draft 组被错误合并。
• tests/v1/attention/test_mla_noncausal.py（模块 非因果测试；类别 test；类型 test-coverage；符号 test_mla_cache_marker_is_promoted_to_group_capability, test_mla_cache_marker_is_a_group_property）: 更新 group capability 与 group property 相关断言，适配 merge 语义从 any() 到组内一致的调整。
• tests/v1/attention/test_mla_backends.py（模块 后端测试；类别 test；类型 test-coverage；符号 test_flashinfer_mla_dspark_dcp_supports_target_only, test_flashinfer_mla_dspark_dcp_supports_target_and_draft）: 更新 DSpark DCP 支持相关断言，验证 target-only 与 target+draft 两组配置在 FlashInfer MLA 下的 capability。

关键符号：BaseSpeculator._build_draft_attn_metadata, DFlashSpeculator._build_draft_attn_metadata, AttentionSpec.merge, FlashInferMLAImpl.forward_mqa, FlashInferMLAMetadataBuilder.supports_non_causal_multi_token_dcp, prepare_dcp_local_seq_lens, test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens, test_flashinfer_mla_forward_uses_native_dcp_api

关键源码片段

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

核心变更点：在共享 draft 元数据构建器中根据 draft 模型当前全局序列长度重算 rank-local DCP 长度，所有 speculator（EAGLE/MTP/DFlash）统一受益。

def _build_draft_attn_metadata(
    self,
    num_reqs: int,
    num_reqs_padded: int,
    num_tokens_padded: int,
    seq_lens_cpu_upper_bound: torch.Tensor,
    step: int,
    num_query_per_req: int | None = None,
    causal: bool | Mapping[int, bool] = False,
    query_start_loc_np: np.ndarray | None = None,
    dcp_local_seq_lens: torch.Tensor | None = None,
) -> dict[str, Any] | None:
    # 前面已构建 query_start_loc_cpu、block_tables、slot_mappings 等
    # 与 DCP 无关的张量，这里从略。
    draft_seq_lens_cpu_upper_bound = torch.zeros(
        num_reqs_padded, dtype=torch.int32, device="cpu"
    )
    torch.add(
        seq_lens_cpu_upper_bound[:num_reqs],
        step,
        out=draft_seq_lens_cpu_upper_bound[:num_reqs],
    )
    draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
​
    if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
        # Draft 步骤在采样与拒绝后会推进或回退自己的全局序列长度，
        # 因此 target 模型遗留的 DCP 本地长度可能已过期，必须基于当前
        # input_buffers.seq_lens 重新计算 rank-local 边界。
        prepare_dcp_local_seq_lens(
            self.input_buffers.dcp_local_seq_lens,
            self.input_buffers.seq_lens,
            num_reqs,
            self.block_tables.cp_size,
            self.block_tables.cp_rank,
            self.block_tables.cp_interleave,
        )
        dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
​
    return build_attn_metadata(
        attn_groups=self.attn_groups,
        num_reqs=num_reqs_padded,
        num_tokens=num_tokens_padded,
        query_start_loc_gpu=self.input_buffers.query_start_loc[: num_reqs_padded + 1],
        query_start_loc_cpu=query_start_loc_cpu,
        max_query_len=max_query_len,
        seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
        dcp_local_seq_lens=(
            None
            if dcp_local_seq_lens is None
            else dcp_local_seq_lens[:num_reqs_padded]
        ),
        max_seq_len=self.draft_max_seq_len,
        block_tables=block_tables,
        slot_mappings=slot_mappings,
        kv_cache_config=self.kv_cache_config,
        causal=causal,
        seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,
    )

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

核心 backend 变更：声明非因果多 token DCP 支持，并在 forward_mqa 中对全局因果边界做几何展开。

# 位于 FlashInferMLAImpl.forward_mqa 的 decode 分支内。
if not attn_metadata.causal:
    # FlashInfer decode 内核没有 causal 标志，对于非因果的 DSpark
    # 草稿块，需要把每个 query 块展开为独立的单 token 解码行。
    query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
    q = q.unsqueeze(1)
    if query_len > 1:
        q = q.repeat(1, query_len, 1, 1)
    # 这里还会对 block_table 与 seq_lens 做 repeat_interleave，
    # 与下方全局因果边界的展开保持相同的几何形状。
​
# DCP 路径必须提供全局因果边界，帮助 native CP 内核恢复
# 跨 rank 的因果掩码信息。
if self.dcp_world_size > 1:
    assert self.cp_kv_cache_interleave_size == 1
    causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens
    assert causal_seqlens_kv_global is not None
    if not attn_metadata.causal and query_len > 1:
        # 非因果块被展开成 query_len 个 decode 行后，全局因果边界
        # 也必须按相同倍数重复，否则 native CP 内核的 batch 几何
        # 会与局部 block table / seq_lens 不匹配。
        causal_seqlens_kv_global = causal_seqlens_kv_global.repeat_interleave(
            query_len
        )
    extra_kwargs.update(
        enable_dcp=True,
        cp_world=self.dcp_world_size,
        cp_rank=self.dcp_rank,
        causal_seqlens_kv_global=causal_seqlens_kv_global,
    )
# 后续调用 trtllm_batch_decode_with_kv_cache_mla 并传入 extra_kwargs。

vllm/v1/kv_cache_interface.py

收紧 AttentionSpec.merge 对 non_causal_multi_token_decode 的一致性要求，避免因果 target 组与非因果 draft 组被错误合并。

@classmethod
def merge(cls, specs: list[Self]) -> Self:
    # 同一 KV cache group 内的所有层必须共享相同的量化方式、
    # tokens_per_state 与 model_version，否则直接报错。
    # 这里省略前面的前置一致性检查。
​
    # non_causal_multi_token_decode 决定该 group 是否以非因果模式
    # 处理多 token 解码。若 group 内同时存在 causal target 层与
    # 非 causal draft 层，旧的 any() 语义会把整个 group 提升为
    # 非因果，导致 target 侧掩码错误。这里改为要求组内完全一致。
    non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
    assert len(non_causal_mtd_set) == 1, (
        "All attention layers in the same KV cache group must agree on "
        "non_causal_multi_token_decode."
    )
​
    merged_spec = cls(
        block_size=specs[0].block_size,
        num_kv_heads=specs[0].num_kv_heads,
        # 其余字段均从 specs 中推导，此处省略。
        non_causal_multi_token_decode=non_causal_mtd_set.pop(),
    )
    # 逐 spec 校验与返回逻辑省略。
    return merged_spec

评论区精华

该 PR 没有产生逐行 review 评论，核心结论来自两位维护者批准：TheEpicDolphin 明确表示“The speculator changes LGTM”，ywang96 直接批准。CI 侧通过 github-actions 触发了 Buildkite CI #86061。PR body 额外披露 AI 辅助用于对比 k3-dev、追踪共享 draft metadata 路径、编写初始 patch 与测试，并承诺提交者将逐行审阅并完成端到端评估。

• Speculator 变更审查 (other): 两组变更均获批准，PR 被合并。
• CI 触发与验证范围 (testing): 已验证单元测试，端到端评估需后续补做。
• AI 辅助披露 (other): 已披露，待人工复核与端到端验证。

风险与影响

• 风险：
  • 缺少端到端验证：PR body 明确说明 Kimi-K3 DSpark+DCP 端到端评估因 GPU 占用未执行，当前仅依赖 12 个 B300 单元测试，几何/数值错误可能在长上下文或真实采样路径中才暴露。
  • KV cache 组合并语义收紧：vllm/v1/kv_cache_interface.py 的 AttentionSpec.merge 从 any() 改为断言组内一致，若现存模型配置中存在因果与非因果层共享 KV cache group 的情况，将直接启动失败。这在语义上正确，但属于行为变化，需排查受影响模型。
  • DCP 长度重算依赖 BlockTables 几何：speculator.py 中 prepare_dcp_local_seq_lens 依赖 block_tables.cp_size/cp_rank/cp_interleave，若这些字段与实际 DCP 配置不同步会引入错误；现有测试只覆盖了 EAGLE 路径。
  • repeat_interleave 放大张量：非因果多 token 路径会把 causal_seqlens_kv_global 按 query_len 重复，增加显存与内核调度开销，对超大 batch 可能敏感。
  • 影响：用户侧：DSpark + FlashInfer MLA + DCP 的组合首次可用，同时 EAGLE、MTP、DFlash 在 DCP 下的 draft 元数据更加准确。系统侧：共享 draft metadata 构建器减少重复代码，后续新增 speculator 可直接继承；KV cache group 语义更严格，可能提前暴露配置错误。团队侧：FlashInfer 后端维护者需要关注 supports_non_causal_multi_token_dcp 能力标记的传播；投机解码模块的 DCP 路径自此统一。
  • 风险标记：缺少端到端 DSpark+DCP 验证, KV cache 组合并语义收紧, 核心注意力路径变更, 依赖 FlashInfer 专属行为

关联脉络

• PR #54012 Enable FlashInfer MLA DCP support for causal decoding (native CP path): PR body 明确说明本 PR 构建于其 causal native FlashInfer MLA DCP 支持之上，本 PR 将能力扩展到非因果 DSpark 草稿路径。
• PR #54102 [SpecDecode] Preserve reject-only residual verifier padding: PR body 说明本 PR 与其相互独立，但同属 speculative decoding + DCP/verifier 功能线；关联 issue 也指向该 PR。