执行摘要

• 一句话：将 DSA top-k 种子锚定至 draft-extend 阶段以提升接受长度
• 推荐动作：值得精读，尤其关注 seed 传递在 eager、CUDA graph、overlap 三种路径下的处理方式。设计决策明确（使用门控开关确保兼容性），展示了如何通过调整缓存锚定点优化推测解码。对于从事推测解码或 MOE 模型性能优化的工程师有参考价值。

功能与动机

PR 描述中指出：'GLM-5.2's MTP (NextN) draft reuses the DSA indexer top-k across draft-decode steps (index_share_for_mtp_iteration) rather than recomputing it each step. Currently the reused top-k is captured on the first draft-decode step, whose query hidden state is the draft model's own step-0 output. That anchor is one step too late — the IndexShare seed should come from the last verified token's (target-derived) hidden state produced during draft-extend.' 此外，zRzRzRzRzRzRzR 在评论中总结：'#29654 fixed the reuse mechanism itself... #29787 then refines where the seed comes from for GLM-5.2 MTP: capture the indexer topk during draft-extend on the last verified token, then reuse it in the next draft loop.'

实现拆解

1. 初始化 DSA index-share 状态（eagle_worker_v2.py）：在 EagleDraftWorker.__init__ 中调用新的 _init_dsa_index_share_state 方法，从模型配置中读取 index_share_for_mtp_iteration 和 index_topk，计算 seed_dsa_topk_from_draft_extend 标志，替换原先简单设置 index_share_for_mtp_iteration 的逻辑。同时预留 dsa_extend_topk_buf 张量用于 eager 模式。

2. 在 draft-extend 前向传播时捕获 top-k 种子（deepseek_nextn.py）：修改 DeepseekV3ForCausalLMNextN.forward，在 decoder 返回 topk_indices 后，检查 forward_mode.is_extend(include_draft_extend_v2=True)，若满足且 dsa_seed_topk_capture 存在，则通过 dsa_seed_topk_select 选择最后一个位置的索引（如果提供），将对应的 top-k 复制到 dsa_seed_topk_capture 中。这样 draft-extend 最后一个验证 token 的 indexer top-k 被保存下来。

3. 在 draft 解码循环中复用种子（eagle_worker_v2.py 的 draft_forward）：当 seed_dsa_topk_from_draft_extend 开启且 spec_info.dsa_topk_indices 不为空时，保持 dsa_topk_indices 不变（复用种子），否则设置为 None 以允许步骤 0 重新计算。同时设置 forward_batch.reuse_dsa_topk_indices = True，让 attention 后端跳过 indexer top-k 计算。

4. CUDA graph 静态缓冲区适配（eagle_draft_cuda_graph_runner.py 和 eagle_draft_extend_cuda_graph_runner.py）：为 draft 和 draft-extend 的 CUDA graph runner 添加 dsa_seed_topk 和 dsa_seed_topk_capture 缓冲区。在 capture_one_shape 中将缓冲区关联到 spec_info，在 execute replay 时从 spec_info 复制到缓冲区（反之亦然）。确保静态图内 seed 正确传递。

5. Overlap schedule 中继（overlap_utils.py）：RelayPayload 新增 dsa_topk_indices 字段，FutureMap 新增 dsa_topk_indices_buf。在 _resolve_spec_extras 和 stash 方法中传递 seed，使得异步 batch 合并场景下种子不丢失。

6. 数据结构重命名与批操作适配（eagle_info.py）：EagleDraftInput 将 mtp_topk_indices 重命名为 dsa_topk_indices，并更新 filter_batch 和 merge_batch。EagleDraftExtendInput 新增 dsa_seed_topk_capture 和 dsa_seed_topk_select。ForwardBatch 中的 reuse_mtp_topk_indices 同步改为 reuse_dsa_topk_indices。

7. 测试 mock 更新：在 attention unit test harness 和单元测试中 mock 新增属性（seed_dsa_topk_from_draft_extend、dsa_index_topk），防止 AttributeError。

关键文件：

• python/sglang/srt/speculative/eagle_worker_v2.py（模块 推测解码；类别 source；类型 core-logic；符号 _init_dsa_index_share_state, _get_dsa_extend_topk_buf）: 核心入口，新增 _init_dsa_index_share_state 方法，修改 draft_forward 种子传递逻辑，是变更的主干。
• python/sglang/srt/models/deepseek_nextn.py（模块 模型层；类别 source；类型 data-contract）: 在该模型的 forward 中添加了 draft-extend 阶段的 seed 捕获逻辑，是种子锚点的产生位置。
• python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py（模块 CUDA图；类别 source；类型 core-logic）: 为 draft-decode 的 CUDA graph 添加了 dsa_seed_topk 缓冲区，确保静态图内种子传递。
• python/sglang/srt/managers/overlap_utils.py（模块 调度器；类别 source；类型 core-logic）: 在 RelayPayload 和 FutureMap 中新增 dsa_topk_indices 字段，支持 overlap 路径下的种子中继。
• python/sglang/srt/speculative/eagle_info.py（模块 推测解码；类别 source；类型 core-logic）: 将 mtp_topk_indices 重命名为 dsa_topk_indices，并更新 filter/merge 方法，同时为 EagleDraftExtendInput 新增 seed 捕获字段。
• python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py（模块 CUDA图；类别 source；类型 core-logic）: 为 draft-extend 的 CUDA graph 添加了 dsa_seed_topk_capture 缓冲区，用于在静态图内捕获种子。

关键符号：_init_dsa_index_share_state, draft_forward, DeepseekV3ForCausalLMNextN.forward, RelayPayload.from_draft_input, FutureMap._resolve_spec_extras, FutureMap.stash, EagleDraftInput.filter_batch, EagleDraftInput.merge_batch, EAGLEDraftCudaGraphRunner.execute, EAGLEDraftExtendCudaGraphRunner.capture_one_shape

关键源码片段

python/sglang/srt/speculative/eagle_worker_v2.py

核心入口，新增 _init_dsa_index_share_state 方法，修改 draft_forward 种子传递逻辑，是变更的主干。

def _init_dsa_index_share_state(self) -> None:
    # 从 draft 模型的 hf_config 中读取 DSA index-share 相关配置
    hf_config = self.draft_runner.model_config.hf_config
    # 是否允许跨 MTP 步骤复用 indexer top-k（要求 topk == 1）
    self.index_share_for_mtp_iteration = (
        getattr(hf_config, "index_share_for_mtp_iteration", False)
        and self.topk == 1
    )
    # indexer 的 top-k 大小（如 128）
    self.dsa_index_topk = getattr(hf_config, "index_topk", None)
    # 启用从 draft-extend 捕获种子（而不是 draft-decode step 0）
    self.seed_dsa_topk_from_draft_extend = (
        self.index_share_for_mtp_iteration and self.dsa_index_topk is not None
    )

在 draft_forward 中的关键改动：

if self.index_share_for_mtp_iteration:
    forward_batch.reuse_dsa_topk_indices = True
    # 如果 seed_dsa_topk_from_draft_extend 开启且已有种子，则保持种子不变
    if not (
        self.seed_dsa_topk_from_draft_extend
        and spec_info.dsa_topk_indices is not None
    ):
        # 否则清除种子，让 step 0 重新计算
        spec_info.dsa_topk_indices = None

python/sglang/srt/models/deepseek_nextn.py

在该模型的 forward 中添加了 draft-extend 阶段的 seed 捕获逻辑，是种子锚点的产生位置。

# 在 DeepseekV3ForCausalLMNextN.forward 的 decoder 调用之后追加：
if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
    # 尝试从 spec_info 中获取 seed capture 缓冲区
    seed_buf = forward_batch.spec_info.dsa_seed_topk_capture
    if seed_buf is not None and topk_indices is not None:
        # 如果提供了 select 索引，则选取特定位置（通常是最后一个 token）
        sel = forward_batch.spec_info.dsa_seed_topk_select
        src = topk_indices if sel is None else topk_indices[sel]
        # 将 last-token 的 indexer top-k 复制到 seed 缓冲区
        seed_buf[: src.shape[0]].copy_(src)

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

为 draft-decode 的 CUDA graph 添加了 dsa_seed_topk 缓冲区，确保静态图内种子传递。

# 在 __init__ 中，创建 dsa_seed_topk 缓冲区
dsa_seed_topk = (
    torch.zeros(
        (self.max_bs, self.eagle_worker.dsa_index_topk),
        dtype=torch.int32,
        device=model_runner.device,
    )
    if self.eagle_worker.seed_dsa_topk_from_draft_extend
    else None
)
​
# 在 execute 方法中，将种子从 spec_info 复制到静态缓冲区
if buffers.dsa_seed_topk is not None:
    seed = forward_batch.spec_info.dsa_topk_indices
    if seed is not None:
        buffers.dsa_seed_topk[:raw_bs].copy_(seed)
    else:
        buffers.dsa_seed_topk[:raw_bs].zero_()

评论区精华

关键讨论：

• zRzRzRzRzRzRzR 评论：'#29654 fixed the reuse mechanism itself... #29787 then refines where the seed comes from for GLM-5.2 MTP: capture the indexer topk during draft-extend on the last verified token, then reuse it in the next draft loop. If the draft-extend seed is unavailable, falling back to draft step 0 computation is still valid.' 这明确了本 PR 的定位以及回退机制。

• JustinTong0323 提供了 torch-profiler 结果（B300, eager 模式），确认 'the reused indexer top-k recompute is actually eliminated in the draft-decode loop'，从性能层面验证了改造效果。

• ch-wan 最终审批通过（LGTM）。

• IndexShare 种子锚点定位 (design): 确认了锚点移动至 draft-extend 是合理的设计，并有回退机制保证兼容性。

• 性能验证 (performance): 实验证实了优化有效性，消除了不必要的重计算。

风险与影响

• 风险：
  1. 核心路径变更风险：draft_forward 是推理关键路径，修改了 topk 种子传递逻辑，可能影响非 GLM-5.2 模型。但已通过 seed_dsa_topk_from_draft_extend 门控（仅当 index_share_for_mtp_iteration 为真且 index_topk 存在时启用），确保兼容性。
  2. 缺少集成测试：仅单元测试通过 mock 验证，未覆盖端到端实际推理场景，可能遗漏运行时细节问题（如缓冲区大小不匹配、异步中继竞态）。
  3. CUDA graph 路径假设：静态图种子传递依赖缓冲区正确关联，capture_one_shape 和 execute 中的 shape 隐式一致，一旦配置变化可能导致静默错误。
  4. Overlap 调度可靠性：种子通过 FutureMap 中继涉及多个异步步骤，代码中已使用 record_stream 确保设备同步，但仍可能引入微妙的时序依赖。
     - 影响：用户影响：仅对启用 index_share_for_mtp_iteration 的 GLM-5.2 MTP 模型生效，长上下文场景接受长度提升（基准测试 OpenHands +0.061）；短上下文和未启用该特性的模型无影响。

系统影响：增加了少量 GPU 内存用于种子缓冲区（每个请求约 index_topk 个 int32），开销可忽略。无性能退化。

团队影响：无新建模块或接口变更，对依赖此代码的其他团队透明。

• 风险标记：核心路径变更（推理关键路径受影响）, 缺少集成测试（仅单元测试 mock）, 兼容性门控（通过 seed_topk_from_extend 开关）, CUDA graph 和 overlap 路径的静态/异步一致性

关联脉络

• PR #29654 [Spec] Fix IndexShare carry loss across per-step forwards: 前身 PR，修复了 index-share 重用机制本身；本 PR 在此基础上进一步优化种子锚点