# PR #29124 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Spec] Unify the overlap stash relay behind a RelayPayload dataclass
- 合并时间：2026-06-25 06:36
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/29124

---

# 执行摘要

- 一句话：统一 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 导入。

```python
@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 开发。