Prhub

#24735 [Spec] Move `accept_tokens` off `EagleDraftInput`; pass via method arg

原始 PR 作者 hnyls2002 合并时间 2026-05-09 14:24 文件变更 7 提交数 10 评论 5 代码增减 +144 / -94

执行摘要

将 accept_tokens 从 EagleDraftInput 移除,改为方法参数传递

在 #24724 中拆分了 verified_id 后,accept_tokens 仍然挂在 EagleDraftInput 上,导致 draft input 承载了 verify output 的职责。本 PR 将其彻底迁移到 EagleVerifyOutput 中,通过方法参数传递,使接口职责更单一。

值得精读,尤其是理解 speculative decode 数据流如何逐步解耦。PR 设计上分离 draft input 与 verify output 的职责,是很好的架构演进方向。

讨论亮点

该 PR 没有实质性的 review 讨论,仅包含 CI 测试触发指令和 bot 回复。

实现拆解

  1. 在 EagleVerifyOutput 中新增 create_idle 静态方法,统一 idle 场景下的输出创建,避免直接构造时传递大量空张量。
  2. 从 EagleDraftInput 中移除 accept_tokens、seq_lens_for_draft_extend、seq_lens_for_draft_extend_cpu、req_pool_indices_for_draft_extend 等字段,转移到 EagleVerifyOutput 中。
  3. 修改 forward_draft_extend_after_decode 和 check_forward_draft_extend_after_decode 方法,使其直接接收 EagleVerifyOutput 参数,从 verify_output 中读取 unfinished_accept_tokens 等信息,不再依赖 batch.spec_info.accept_tokens。
  4. 在 verify 方法的返回处,将 draft_input 重命名为 next_draft_input,并调整 EagleVerifyOutput 的构造方式,同时将用于 extend 的序列信息直接放入 verify_output。
  5. 同步更新 frozen_kv_mtp、dflash、ngram 等 info 类,移除遗留的 self.last_loc 赋值,并适配新的接口。
文件 模块 状态 重要度
python/sglang/srt/speculative/eagle_info.py 推测解码 modified 7.62
python/sglang/srt/speculative/eagle_worker.py 推测解码 modified 7.15
python/sglang/srt/speculative/multi_layer_eagle_worker.py 推测解码 modified 7.15
python/sglang/srt/speculative/frozen_kv_mtp_worker.py 推测解码 modified 6.84
python/sglang/srt/speculative/frozen_kv_mtp_info.py 推测解码 modified 4.82
python/sglang/srt/speculative/dflash_info.py 推测解码 modified 4.58
python/sglang/srt/speculative/ngram_info.py 推测解码 modified 4.58

关键符号

EagleVerifyOutput.create_idle EagleVerifyOutput.verify forward_draft_extend_after_decode check_forward_draft_extend_after_decode

关键源码片段

python/sglang/srt/speculative/eagle_info.py data-contract

核心数据结构变更,移除了 EagleDraftInput.accept_tokens 等字段,新增 EagleVerifyOutput.create_idle 工厂方法,修改 verify 方法返回结构。

class EagleInfo:
    def verify(self, batch: ScheduleBatch, logits_output: ...) -> EagleVerifyOutput:
        if batch.forward_mode.is_idle():
            # 空闲模式:创建空闲的 draft input,然后通过工厂方法构建 verify output
            next_draft_input = EagleDraftInput.create_idle_input(
                device=batch.device,
                hidden_size=batch.model_config.spec_hidden_size,
                dtype=batch.model_config.dtype,
                topk=self.topk,
                capture_hidden_mode=CaptureHiddenMode.LAST,
            )
            return EagleVerifyOutput.create_idle(
                next_draft_input=next_draft_input,
                logits_output=logits_output,
                device=batch.device,
                spec_steps=self.spec_steps,
            )
        # 非空闲分支省略 ...

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

核心数据结构 EagleDraftInput 和 EagleVerifyOutput 的字段被重命名和移除,所有调用点(EagleWorker、MultiLayerEagleWorker、FrozenKVMTPWorker)已同步修改,但若存在未覆盖的第三方自定义 speculative 算法或分支,可能导致运行时属性缺失错误。另外,unfinished_accept_tokens 与原有 accept_tokens 的语义区分(子集 vs 全集)需要严格保证一致性,否则会影响 extend 前的数据准备。

影响范围限于 speculative decoding 模块的 worker 和 info 类,用户无感知。内部数据流更清晰,为后续支持多 draft 或更复杂的 verify 逻辑打下基础。建议合并后运行 commited 测试套件(如 test_eagle_dp_attention.py、test_deepseek_v3_mtp.py)确保回归覆盖。

核心数据结构变更 缺少直接测试覆盖 多 worker 同步风险

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论