# PR #26972 完整报告

- 仓库：`sgl-project/sglang`
- 标题：Spec v2 tree drafting (topk>1) with page_size>1
- 合并时间：2026-06-07 03:00
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/26972

---

# 执行摘要

- 一句话：扩展 EAGLE spec v2 树推理支持 page>1+topk>1
- 推荐动作：建议所有涉及 speculative decoding 的开发者和 reviewer 精读。关键设计决策包括：孔状布局理由、前缀复制方案、行宽保护与 failure 模式选择。值得关注的设计模式：用 always-on CPU 断言代替难诊断的 GPU 错误，是防守型编程的良好范例。

# 功能与动机

PR body 指出，当 page_size>1 且 topk>1 时，draft 预留的页对齐足迹（2 * get_alloc_len_per_decode）远超默认的 4+num_draft_tokens 头空间，导致 free 侧 KV 泄漏（pool memory leak detected）和 access 侧 OOB（CUDA gather 断言）。本 PR 旨在解锁 page>1+topk>1 组合在 spec v2 中的正确使用，同时通过行宽调整和不变性检查根除这两种失败模式。

# 实现拆解

1. **布局与复制**：在 `eagle_info_v2.py` 中新增 `duplicate_prefix_tail_to_draft_branches` 函数，将前缀的部分尾页复制到每个分支（branch b≥1）的第一个页孔中，确保注意力整页读取一致。
2. **分配长度计算**：在 `managers/utils.py` 的 `get_alloc_len_per_decode` 中，为 `page_size>1+topk>1` 实现最坏情况分配长度：`num_new_pages_per_topk * page_size * topk`，替换之前的 `NotImplementedError`。
3. **池大小扩宽**：在 `model_runner_kv_cache_mixin.py` 的 `_init_pools` 中，当满足 `speculative_algorithm非None且page_size>1且topk>1` 时，将 `extra_max_context_len` 提升为 `max(4+num_draft_tokens, 2*get_alloc_len_per_decode)`，从而加宽 `req_to_token` 行，容纳 draft 的孔状占用。
4. **运行时保护**：在 `eagle_info_v2.py` 的 `prepare_for_decode` 中添加始终开启的 CPU 端断言，检查本批次最大分配长度不超过行宽，以清晰错误代替静默损坏。
5. **路由与后端门控**：在 `arg_groups/speculative_hook.py` 中，移除 `page_size>1` 强制回退 spec v1 的限制，改为仅针对 mamba 状态模型强制回退；同时添加后端白名单（flashinfer、fa3、triton），对不支持的 attention backend 报错。
6. **测试配套**：新增 `test/registered/spec/eagle/test_spec_eagle_topk_page.py`，包含 `TestEagle3Page64Topk8`（spec v2）和 `TestEagleLlama2Page4Topk8`（spec v1）两个测试用例。修改 `test_spec_eagle_page.py` 移除重复符号。

关键文件：
- `python/sglang/srt/speculative/eagle_info_v2.py`（模块 推测解码；类别 source；类型 core-logic；符号 duplicate_prefix_tail_to_draft_branches, prepare_for_decode）: 核心实现：新增 duplicate_prefix_tail_to_draft_branches 函数实现前缀尾页复制；在 prepare_for_decode 中添加行宽不变性检查。
- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`（模块 内存池；类别 source；类型 data-contract）: 池初始化：在 _init_pools 中加宽 req_to_token 行宽以容纳 draft 孔状占用，避免 KV 泄漏和 OOB。
- `test/registered/spec/eagle/test_spec_eagle_topk_page.py`（模块 测试；类别 test；类型 test-coverage；符号 TestEagle3Page64Topk8, TestEagleLlama2Page4Topk8）: 新增测试用例覆盖 EAGLE3（page64+topk8）和 Llama-2（page4+topk8），确保功能正确性。
- `python/sglang/srt/arg_groups/speculative_hook.py`（模块 参数配置；类别 source；类型 core-logic；符号 _handle_eagle_family）: 路由修改：解除 page_size>1 对 spec v2 的限制，添加后端白名单检查。
- `python/sglang/srt/managers/utils.py`（模块 工具函数；类别 source；类型 core-logic；符号 get_alloc_len_per_decode）: 分配计算：实现 page_size>1+topk>1 的 get_alloc_len_per_decode，此前为 NotImplementedError。

关键符号：duplicate_prefix_tail_to_draft_branches, get_alloc_len_per_decode, _init_pools, _handle_eagle_family, prepare_for_decode

## 关键源码片段

### `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`

池初始化：在 _init_pools 中加宽 req_to_token 行宽以容纳 draft 孔状占用，避免 KV 泄漏和 OOB。

```python
# Inside _init_pools, after the base extra_max_context_len based on num_draft_tokens:

# page>1 + topk>1 reserves a holey draft footprint (2 * get_alloc_len_per_decode
# = topk * num_new_pages * page) far beyond the default num_draft_tokens
# headroom; widen the row to hold it, else free leaks KV and the holey gather OOBs.
if (
    self.server_args.speculative_algorithm is not None
    and self.server_args.page_size > 1
    and (self.server_args.speculative_eagle_topk or 1) > 1
):
    from sglang.srt.managers.utils import get_alloc_len_per_decode

    extra_max_context_len = max(
        extra_max_context_len,
        2 * get_alloc_len_per_decode(self.server_args),
    )

```

# 评论区精华

开发者在 PR body 中详细描述了 bug 的两个表现（free-side leak 和 access-side OOB），并提供了两个 repro gist（page512 multi-turn 和 page256 long prompt）以及用 `SGLANG_DEBUG_REVERT_PR=26972` 验证修复的 A-B 测试。后续关联 PR #27338（flashinfer cuda-graph kv_indices）和 #27360（fa3 expand page_table）继续修复其他 backends 上的同类问题。无额外 review 讨论。

- 暂无高价值评论线程

# 风险与影响

- 风险：
 1. **内存占用增加**：`_init_pools` 中加宽行宽可能增加 `req_to_token` 池内存，尤其在高 `page_size` 或 `topk` 下；但由于仅在启用 page>1+topk>1 时生效，风险可控。
 2. **后端兼容性**：仅 flashinfer、fa3、triton 后端支持 page>1+topk>1，若用户尝试其他后端（如 flashmla、trtllm_mla）会报错，需确保文档清晰。
 3. **GPU kernel 调用**：`duplicate_prefix_tail_to_draft_branches` 调用 `token_to_kv_pool.move_kv_cache`，这是一个 GPU 操作，可能增加延迟，但仅每轮 decode 一次，影响有限。
 4. **不变性断言性能**：`prepare_for_decode` 中的断言涉及 CPU-GPU 同步？实际使用 CPU max 和 shape，开销低，但需注意不会引入阻塞。
 - 影响：**用户影响**：使用 EAGLE 推测解码且 page_size>1 的用户现在可以启用 topk>1（tree drafting），获得更高推测接受率；但必须使用支持的后端。之前该组合要么回退到 spec v1（实际上之前代码导致 prepare_for_decode 报 NotImplementedError），要么崩溃。现在正常运作。
**系统影响**：`req_to_token` 池增大可能略微增加显存占用，但仅在开启时。
**团队影响**：规范解码的维护复杂度增加，需要确保其他 backend（如 flashmla）也能支持；新增的 pr_fix_toggle 注册方便调试和回滚。

- 风险标记：内存占用增加 , 后端兼容限制 , GPU 同步开销

# 关联脉络

- PR #26866 Spec v2 tree drafting (topk>1): 本 PR 栈在该 PR 之上，扩展了其 page_size>1 支持。
- PR #27338 Fix EAGLE draft CUDA-graph `kv_indices` under-allocation for `topk > 1`: 关联修复：flashinfer 后端中 draft kv_indices 缓冲区缺少 topk 因子导致 OOB，与本 PR 共同解决 page>1+topk>1 问题。
- PR #27360 Fix fa3 EAGLE draft-decode expand page_table scatter OOB for topk>1 + page_size>1: 关联修复：fa3 后端中 expand page_table scatter OOB，与本 PR 一起确保 page>1+topk>1 在主要后端上正确工作。