执行摘要

• 一句话：支持 FlashMLA 使用 spec decoding v2
• 推荐动作：值得精读，展示了如何为特定 attention 后端添加 spec v2 支持，涉及调度模式匹配和条件分支的技巧。

功能与动机

关联 Issue #24637，当前 FlashMLA 的 CI 测试强制回退到 spec v1（设置 SGLANG_ENABLE_SPEC_V2=False），需要让 FlashMLA 支持 spec v2 并在 CI 中覆盖该路径。

实现拆解

1. 修改 flashmla_backend.py：在 forward_extend 方法中将判断条件从 （EXTEND or DRAFT_EXTEND） 扩展为 in (EXTEND, DRAFT_EXTEND, DRAFT_EXTEND_V2)，使得 spec v2 的 draft extend 阶段复用基类的 prefill 逻辑（使用 FlashInfer MLA 后端），而不是进入 FlashMLA 的 target-verify 路径，从而正确匹配 metadata 初始化。

2. 修改 deepseek_v2.py：在 dispatch_attn_forward_method 方法中将 forward_batch.forward_mode.is_draft_extend() 改为 is_draft_extend(include_v2=True)，确保 spec v2 的 draft extend 阶段也能根据 speculative_attention_mode 正确选择 attention 后端。

3. 修改 test_flashmla.py：移除 envs.SGLANG_ENABLE_SPEC_V2.override(False) 的强制 v1 设置，使测试使用默认的 spec v2，同时移除不再需要的 envs 导入。这使得 CI 能够自动验证 spec v2 路径。

关键文件：

• python/sglang/srt/layers/attention/flashmla_backend.py（模块 FlashMLA 后端；类别 source；类型 core-logic）: 核心逻辑修改：在 forward_extend 中增加 DRAFT_EXTEND_V2 分支，使 spec v2 draft extend 阶段走预填充路径。
• test/registered/mla/test_flashmla.py（模块 测试；类别 test；类型 test-coverage）: CI 测试配置修改：移除强制 spec v1 的环境变量覆盖，使测试默认运行 spec v2 路径。
• python/sglang/srt/models/deepseek_v2.py（模块 DeepSeek 模型；类别 source；类型 data-contract）: 注意力分发逻辑修改：使 is_draft_extend 包含 v2，确保 spec v2 模式也能正确选择 attention 后端。

关键符号：FlashMLABackend.forward_extend, DeepseekV2Attention.dispatch_attn_forward_method

关键源码片段

python/sglang/srt/layers/attention/flashmla_backend.py

核心逻辑修改：在 forward_extend 中增加 DRAFT_EXTEND_V2 分支，使 spec v2 draft extend 阶段走预填充路径。

    def forward_extend(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer: RadixAttention,
        forward_batch: ForwardBatch,
        save_kv_cache: bool = True,
    ):
        # 关键修改：将 DRAFT_EXTEND_V2 加入条件，使 spec v2 的 draft extend
        # 复用基类 forward_extend（使用 FlashInfer MLA），而不是进入下方的
        # FlashMLA target-verify 路径，确保 metadata 正确初始化
        if forward_batch.forward_mode in (
            ForwardMode.EXTEND,
            ForwardMode.DRAFT_EXTEND,
            ForwardMode.DRAFT_EXTEND_V2, # 新增 spec v2 支持
        ):
            return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache)
        else:
            # 原有 FlashMLA 解码 / 校验路径，保持不变
            cache_loc = forward_batch.out_cache_loc
            # ...（后续代码）

评论区精华

Reviewer Qiaolin-Yu 要求提供 profiling 结果以证明操作重叠正确，作者表示缺少 H100 环境；zhendonghua 解释了 flashmla_backend.py 修改的作用；最终 Qiaolin-Yu 在 zhendonghua 验证后批准合并。

• 需提供 profiling 结果验证 overlap 正确性 (correctness): 验证通过，无需额外 profiling。
• flashmla_backend.py 修改解读 (design): 这是预期行为，与 metadata 初始化一致。

风险与影响

• 风险：风险较低。改动集中在三个文件，且为新增分支而非修改现有逻辑。主要风险在于 spec v2 路径之前未在 CI 中覆盖，可能在某些配置下存在隐藏问题；但整体影响范围有限，且经 reviewer 验证。
• 影响：对使用 FlashMLA 后端的用户（主要是 DeepSeek 系列模型）可启用 spec v2，可能带来性能提升或功能兼容性。对非 FlashMLA 后端无影响。CI 测试覆盖面增加，提高了回归安全性。团队维护成本极低。
• 风险标记：核心路径变更, 缺少性能验证

关联脉络

• 暂无明显关联 PR