执行摘要

• 一句话：修复多层级EAGLE推测解码中预填充和解码阶段token池引用不一致的问题。
• 推荐动作：该PR值得快速浏览，重点关注token池引用一致性的设计决策。虽然变更小，但揭示了在多层推测解码中管理状态（如token池）的常见陷阱，对于理解EAGLE架构和避免类似bug有借鉴意义。

功能与动机

PR作者在body中提出疑问：“我认为这是正确的。token_to_kv_pool和attn_backend也需要更新吗？”，表明作者在审查代码时发现了token池引用可能不一致的问题。review评论进一步确认了这个问题：gemini-code-assist[bot]指出“预填充逻辑中将token池分配给worker实例而不是batch对象，这会导致前向传递中token池引用不正确”，并建议“应该更新forward_batch.req_to_token_pool以保持与解码逻辑的一致性”。

实现拆解

1. 识别问题：在multi_layer_eagle_worker_v2.py的_draft_extend_for_prefill函数中，推测解码循环的每个步骤应使用对应draft_runner_list[step]的req_to_token_pool，但代码未设置forward_batch.req_to_token_pool，导致所有步骤都错误地使用第一个draft runner的token池。
2. 修复预填充阶段：在_draft_extend_for_prefill函数的循环开始处（第397-399行），添加forward_batch.req_to_token_pool = self.draft_runner_list[step].req_to_token_pool，确保每个步骤使用正确的token池。
3. 修复解码阶段：在_draft_extend_for_decode函数的非CUDA Graph路径中（第504-506行），同样添加forward_batch.req_to_token_pool = self.draft_runner_list[step].req_to_token_pool，使解码逻辑与预填充保持一致。
4. 验证修复：通过CI测试（如test_mimo_models.py和test_step3p5_flash_chain_mtp.py）确保修复不引入回归，且所有相关测试通过。

关键文件：

• python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py（模块 推测解码；类别 source；类型 core-logic；符号 _draft_extend_for_prefill, _draft_extend_for_decode）: 这是唯一变更的文件，包含多层级EAGLE推测解码的核心逻辑，修复了预填充和解码阶段token池引用不一致的bug。

关键符号：_draft_extend_for_prefill, _draft_extend_for_decode

关键源码片段

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

这是唯一变更的文件，包含多层级EAGLE推测解码的核心逻辑，修复了预填充和解码阶段token池引用不一致的bug。

# 在 _draft_extend_for_prefill 函数中，修复预填充阶段的 token 池引用
for step in range(self.speculative_num_steps):
    # 关键修复：将当前步骤 draft runner 的 token 池赋值给 forward_batch，确保引用正确
    forward_batch.req_to_token_pool = self.draft_runner_list[step].req_to_token_pool
    output: ModelRunnerOutput = self.draft_runner_list[step].forward(forward_batch)
    # ... 后续逻辑保持不变
​
# 在 _draft_extend_for_decode 函数中，修复解码阶段的 token 池引用（非 CUDA Graph 路径）
else:
    # 关键修复：同样更新 forward_batch 的 token 池引用，保持与预填充逻辑一致
    forward_batch.req_to_token_pool = self.draft_runner_list[step].req_to_token_pool
    draft_logits_output = self.draft_runner_list[step].forward(
        forward_batch, skip_attn_backend_init=True
    )
    # ... 后续逻辑保持不变

评论区精华

review中只有一个关键讨论：gemini-code-assist[bot]指出预填充阶段存在token池引用错误，并建议修复。讨论结论是采纳建议，更新forward_batch.req_to_token_pool以保持一致性。未解决疑虑：作者在PR body中询问token_to_kv_pool和attn_backend是否需要更新，但review和后续讨论未涉及，可能留待未来处理。

• token池引用不一致的修复 (correctness): 采纳建议，在预填充和解码阶段都添加forward_batch.req_to_token_pool赋值，确保token池引用正确。

风险与影响

• 风险：技术风险：

  • 回归风险低：修复仅涉及token池引用的一致性，不改变核心算法逻辑，且通过CI测试验证。
  • 性能影响可忽略：添加的赋值操作开销极小，不影响整体性能。
  • 兼容性无影响：修复不改变接口或行为，保持向后兼容。
    具体风险点：如果token_to_kv_pool或attn_backend确实需要类似更新但未处理，可能导致隐藏bug，但当前测试通过表明风险可控。

• 影响：影响范围：

  • 用户影响：修复潜在bug，提升多层级EAGLE推测解码的稳定性和正确性，用户无感知但受益于更可靠的系统。
  • 系统影响：确保推测解码中token池引用正确，避免因引用错误导致的计算异常或性能下降。
  • 团队影响：代码变更小，易于理解和维护，为后续类似修复提供参考。
    影响程度：低到中度，修复核心路径但逻辑简单，影响局限于推测解码模块。

• 风险标记：核心路径变更, 状态管理一致性

关联脉络

• PR #22832 [sgl] fix incorrect behavior in cuda graph draft extend: 修改了同一文件（multi_layer_eagle_worker_v2.py），涉及CUDA Graph推测解码的bug修复，与本PR的推测解码逻辑相关。
• PR #22088 [sgl] add support for weight update function in spedec: 修改了同一文件（multi_layer_eagle_worker_v2.py）和推测解码相关文件，涉及EAGLE推测解码的功能增强，与本PR同属推测解码模块。