执行摘要

• 一句话：统一 FutureMap.stash 为 RelayPayload 数据类
• 推荐动作：值得精读，因为引入了统一数据类设计，消除了类型分支，是清理调度层技术债务的典范。同时展示了如何处理 ngram 等特殊路径。关注点是测试覆盖尚缺，未来合并类似重构时应同步新增测试。

功能与动机

PR Body 中明确指出：FutureMap.stash 原本接受 Union[torch.Tensor, EagleDraftInput]，内部通过类型分派处理不同 payload。替代为单一 RelayPayload 数据类可消除类型开关（解决了一个 FIXME），并为统一解码结果合约奠定基础。同时修复了 ngram 路径因缺少 bonus_tokens 导致的崩溃。

实现拆解

1. 定义 RelayPayload 数据类（overlap_utils.py）：包含必选字段 bonus_tokens（用于非 spec 路径）和可选字段（topk_p, topk_index, hidden_states, draft_probs），并提供 from_draft_input 工厂方法从 EagleDraftInput 构造。
2. 修改 FutureMap 核心方法（overlap_utils.py）：stash 方法签名从 Union[torch.Tensor, EagleDraftInput] 改为 RelayPayload；_lazy_init_forward_buf 相应地接受 RelayPayload 而非 EagleDraftInput。
3. 更新所有调用点：在 scheduler.py、scheduler_pp_mixin.py、decode_schedule_batch_mixin.py、eagle_disaggregation.py 中，将对 stash 的调用均改为构造 RelayPayload 实例。
4. 推迟 spec_info 导入（overlap_utils.py）：将 SpeculativeAlgorithm 等 spec 相关导入推迟到函数内部，保持模块级依赖叶化，避免循环导入。
5. 修复 ngram 兼容性（scheduler.py）：在 stash 调用前增加 is_ngram() 检查，跳过 relay payload 投影，因为 ngram 的 draft input 没有 bonus_tokens。

关键文件：

• python/sglang/srt/managers/overlap_utils.py（模块 调度层；类别 source；类型 core-logic；符号 RelayPayload, from_draft_input, _lazy_init_forward_buf, stash）: 核心变更文件：定义 RelayPayload 数据类，修改 FutureMap.stash 和 _lazy_init_forward_buf 的签名，推迟 spec 导入。
• python/sglang/srt/managers/scheduler.py（模块 调度层；类别 source；类型 core-logic）: 主要调用文件：多处 stash 调用点被更新为 RelayPayload；添加 ngram 兼容检查并修复了崩溃。
• python/sglang/srt/managers/scheduler_pp_mixin.py（模块 调度层；类别 source；类型 dependency-wiring）: PP 调度混合模块：调整 _pp_prep_batch_result 中的 stash 调用为 RelayPayload。
• python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py（模块 分离调度；类别 source；类型 dependency-wiring）: 分离式解码混入：调整 process_prebuilt 中非 spec 路径的 stash 调用为 RelayPayload。
• python/sglang/srt/speculative/eagle_disaggregation.py（模块 推测解码；类别 source；类型 dependency-wiring）: Eagle 分离模块：调整 build_eagle_disagg_draft_input 中 spec 路径的 stash 调用为 RelayPayload。

关键符号：RelayPayload, RelayPayload.from_draft_input, FutureMap._lazy_init_forward_buf, FutureMap.stash, build_eagle_disagg_draft_input

关键源码片段

python/sglang/srt/managers/overlap_utils.py

核心变更文件：定义 RelayPayload 数据类，修改 FutureMap.stash 和 _lazy_init_forward_buf 的签名，推迟 spec 导入。

@dataclass
class RelayPayload:
    """Per-iteration stash payload for the FutureMap bufs. Non-spec fills only
    `bonus_tokens`; which spec extras get relayed is decided by
    `FutureMap.spec_algo`, not by this payload's shape."""
​
    bonus_tokens: torch.Tensor
    topk_p: Optional[torch.Tensor] = None
    topk_index: Optional[torch.Tensor] = None
    hidden_states: Optional[torch.Tensor] = None
    draft_probs: Optional[torch.Tensor] = None
​
    @classmethod
    def from_draft_input(cls, draft_input: EagleDraftInput) -> "RelayPayload":
        """Factory: extract fields from EagleDraftInput. draft_probs may not exist on all subclasses."""
        return cls(
            bonus_tokens=draft_input.bonus_tokens,
            topk_p=draft_input.topk_p,
            topk_index=draft_input.topk_index,
            hidden_states=draft_input.hidden_states,
            draft_probs=getattr(draft_input, "draft_probs", None),
        )
​
# In FutureMap._lazy_init_forward_buf (signature changed):
def _lazy_init_forward_buf(self, payload: RelayPayload):
    # Local import: keep this module leaf-dependency free.
    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
    # ... remaining init uses payload attributes

评论区精华

PR 未产生 review 讨论，但 commit 历史显示了一个关键修复：在运行 ngram 模式时，直接构造 RelayPayload 会因缺少 bonus_tokens 导致崩溃。作者在最后一个 commit 中通过 is_ngram() 检查跳过了相关 stash 调用。

• ngram 兼容性：stash 不允许传入 ngram draft input (correctness): 在 stash 前添加 is_ngram() 检查，跳过 relay payload 投影。

风险与影响

• 风险：核心调度和推测解码的重叠路径被重构，调用点分散在 5 个文件中，如果遗漏某条调用路径或错误构造 RelayPayload，可能导致 relay 数据错误。特别是 ngram 路径的修复尚未有测试覆盖，可能存在类似的边缘情况。另外，RelayPayload 中的可选字段允许 None，消费者需确保健壮性。
• 影响：仅影响启用 overlap 模式且使用 speculative decoding 的调度路径。非 overlap、非 spec 路径完全不受影响。用户功能无可见变化，行为保持。为未来的解码结果统一做铺垫。
• 风险标记：核心路径变更, 缺少测试覆盖, ngram回归已修复

关联脉络

• PR #29122 Unknown (stacked PR): PR body 声明 'Stacks on #29122'，本 PR 基于 #29122 开发。