# PR #24640 完整报告

- 仓库：`sgl-project/sglang`
- 标题：Support spec v2 for FlashMLA speculative decoding
- 合并时间：2026-05-20 06:23
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/24640

---

# 执行摘要

- 一句话：支持 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 阶段走预填充路径。

```python
    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