Spec v2 tree drafting (topk>1) with page_size>1 (#26972)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -285,6 +285,8 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
             "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
         )
 
+    # Spec v2 tree drafting supports topk > 1 with page_size == 1 and page_size > 1
+    # (the latter via partial-page duplication; backend-gated below).
     spec_v1_reason = None
     # mamba / linear-attn state models only support topk == 1 on spec v2.
     # mamba2_cache_params exists iff the config carries such state; check the
@@ -294,16 +296,15 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
     if (
         server_args.speculative_eagle_topk is not None
         and server_args.speculative_eagle_topk > 1
-        and (server_args.page_size > 1 or is_mamba_state_model)
+        and is_mamba_state_model
         and not server_args.disable_overlap_schedule
     ):
-        # Spec v2 topk > 1 only supports page_size == 1 on non-mamba models;
+        # Spec v2 topk > 1 is not supported for mamba/linear-attn state models
-        # page_size > 1 (partial-page dup) isn't ported to v2 yet -> fall back to v1.
+        # (only topk == 1); fall back to v1 for those. page_size > 1 is supported
+        # on v2 (partial-page duplication), so it no longer forces v1.
         server_args.disable_overlap_schedule = True
         spec_v1_reason = (
             "spec v2 topk > 1 is not supported for mamba/linear-attn models"
-            if is_mamba_state_model
-            else "spec v2 topk > 1 currently requires page_size == 1"
         )
     elif (
         not envs.SGLANG_ENABLE_SPEC_V2.get()
@@ -389,13 +390,19 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
         )
         server_args.speculative_num_draft_tokens = server_args.speculative_num_steps + 1
 
+    # topk > 1 + page_size > 1 needs the two-pass cascade draft-decode (shared prefix
+    # pass + per-branch expand pass with prefix-tail dup). Only these backends implement
+    # it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject.
+    _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton")
     if (
         server_args.speculative_eagle_topk > 1
         and server_args.page_size > 1
-        and server_args.attention_backend not in ["flashinfer", "fa3"]
+        and server_args.attention_backend not in _PAGE_TREE_SPEC_BACKENDS
     ):
         raise ValueError(
-            "speculative_eagle_topk > 1 with page_size > 1 is unstable and produces incorrect results for paged attention backends. This combination is only supported for the 'flashinfer' backend."
+            f"speculative_eagle_topk > 1 with page_size > 1 is only supported on "
+            f"{_PAGE_TREE_SPEC_BACKENDS}; got attention_backend="
+            f"{server_args.attention_backend!r}. Use page_size == 1 or one of those backends."
         )
 
 

python/sglang/srt/debug_utils/pr_fix_toggle.py

@@ -87,11 +87,32 @@ patches:
 """
 
 
+_PR_REVERT_YAML_26972 = """
+patches:
+  - target: sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools
+    edits:
+      - match: |
+          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),
+              )
+        replacement: ""
+"""
+
+
 _PR_FIX_REVERT_YAML: Dict[int, str] = {
     25015: _PR_REVERT_YAML_25015,
     26329: _PR_REVERT_YAML_26329,
     27338: _PR_REVERT_YAML_27338,
     27360: _PR_REVERT_YAML_27360,
+    26972: _PR_REVERT_YAML_26972,
 }
 
 

python/sglang/srt/managers/scheduler.py

@@ -419,6 +419,10 @@ class Scheduler(
         # Init mamba backend
         self.init_mamba_backend()
 
+        # Must precede init_model_worker: revert targets like _init_pools run during it,
+        # so patching them afterwards is a no-op.
+        maybe_revert_pr_fix()
+
         # Launch a model worker and draft model worker if using speculative decoding
         self.init_model_worker()
 
@@ -564,8 +568,6 @@ class Scheduler(
 
         self.init_batch_result_processor()
 
-        maybe_revert_pr_fix()
-
         self.is_initializing = False
 
     def init_zbal_on_npu(self):

python/sglang/srt/managers/utils.py

@@ -246,9 +246,14 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
     if page_size == 1 or spec_topk == 1:
         return max(spec_steps * spec_topk, spec_tokens)
     else:
-        raise NotImplementedError(
+        # page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree
-            "get_alloc_len_per_decode not implemented for page_size > 1 and spec_topk > 1"
+        # footprint. Per topk branch needs ceil((last_page_len + num_steps) / page)
-        )
+        # pages; the partial tail page can be up to page_size - 1, and each branch
+        # gets its own (duplicated) copy -- so reserve for all topk branches.
+        num_new_pages_per_topk = (
+            (page_size - 1) + spec_steps + page_size - 1
+        ) // page_size
+        return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens)
 
 
 @dataclass

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -300,6 +300,21 @@ class ModelRunnerKVCacheMixin:
             if max_spec_draft_tokens is not None:
                 extra_max_context_len += max_spec_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),
+                )
+
             if self.server_args.disaggregation_mode == "decode":
                 from sglang.srt.disaggregation.decode import (
                     DecodeReqToTokenPool,

python/sglang/srt/speculative/eagle_info_v2.py

@@ -74,6 +74,53 @@ if is_cuda() or is_musa():
     )
 
 
+def duplicate_prefix_tail_to_draft_branches(
+    token_to_kv_pool,
+    rows: torch.Tensor,
+    prefix_base: torch.Tensor,
+    last_page: torch.Tensor,
+    num_new_pages: torch.Tensor,
+    topk: int,
+    page_size: int,
+) -> None:
+    """Copy the prefix partial-tail page into each branch's first-page holes (page>1 + topk>1).
+
+    The draft-decode expand pass reads each branch's own draft page by block id
+    (cache_loc // page_size), so branch b>=1's hole slots [0, last_page) must hold the
+    real prefix tail (branch 0's first page already is it). Mirrors V1 #7725.
+    """
+    if topk <= 1:
+        return
+    bs = rows.shape[0]
+    page_off = torch.arange(page_size, device=rows.device, dtype=torch.int64)
+    branches = torch.arange(1, topk, device=rows.device, dtype=torch.int64).view(
+        1, topk - 1, 1
+    )
+    # Source: the prefix tail page [prefix_base, prefix_base + page_size), one per branch.
+    src_pos = (prefix_base.view(bs, 1, 1) + page_off.view(1, 1, page_size)).expand(
+        bs, topk - 1, page_size
+    )
+    # Target: branch b's first page [prefix_base + b*num_new_pages*page, + page_size).
+    tgt_pos = (
+        prefix_base.view(bs, 1, 1)
+        + branches * (num_new_pages.view(bs, 1, 1) * page_size)
+        + page_off.view(1, 1, page_size)
+    )
+    # Only [0, last_page) holds real prefix KV; [last_page, page_size) are the branch's
+    # own draft slots and must not be overwritten.
+    vmask = (page_off.view(1, 1, page_size) < last_page.view(bs, 1, 1)).expand(
+        bs, topk - 1, page_size
+    )
+    src_slots = torch.gather(rows, 1, src_pos.reshape(bs, -1)).reshape(
+        bs, topk - 1, page_size
+    )[vmask]
+    tgt_slots = torch.gather(rows, 1, tgt_pos.reshape(bs, -1)).reshape(
+        bs, topk - 1, page_size
+    )[vmask]
+    if src_slots.numel() > 0:
+        token_to_kv_pool.move_kv_cache(tgt_slots, src_slots)
+
+
 @dataclass
 class EagleDraftInputV2Mixin:
     def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
@@ -128,6 +175,21 @@ class EagleDraftInputV2Mixin:
         cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")
         nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu")
 
+        # Fail fast if the page>1 + topk>1 draft over-allocation
+        # (2 * get_alloc_len_per_decode) outgrows the req_to_token row: the write below
+        # would OOB and free would leak KV. The row is widened to hold it in _init_pools
+        # (PR #26972); fail here with a clear error, not on a later cryptic CUDA assert.
+        from sglang.srt.server_args import get_global_server_args
+
+        if page_size > 1 and (get_global_server_args().speculative_eagle_topk or 1) > 1:
+            max_alloc_len = int(nxt_kv_lens_cpu.max())
+            row_width = batch.req_to_token_pool.req_to_token.shape[1]
+            assert max_alloc_len <= row_width, (
+                f"spec v2 page>1 topk>1 draft over-allocation ({max_alloc_len}) exceeds "
+                f"req_to_token row width ({row_width}); page_size={page_size}. Widen the "
+                f"row to hold committed + 2 * get_alloc_len_per_decode (PR #26972)."
+            )
+
         # non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR
         # barrier has chained to the prev forward -> host stalls a full forward.
         cur_kv_lens_device = cur_kv_lens_cpu.to(device=batch.device, non_blocking=True)
@@ -171,22 +233,65 @@ class EagleDraftInputV2Mixin:
         if not batch.forward_mode.is_idle():
             bs = len(batch.seq_lens)
 
-            # Assign cache locations
+            # Assign cache locations (draft-write targets).
-            batch.out_cache_loc = torch.empty(
+            page_size = batch.token_to_kv_pool_allocator.page_size
-                (bs * topk * num_steps,),
+            if page_size == 1 or topk == 1:
-                dtype=torch.int64,
+                batch.out_cache_loc = torch.empty(
-                device=batch.device,
+                    (bs * topk * num_steps,),
-            )
+                    dtype=torch.int64,
-            # FIXME(lsyin): align with the default code path
+                    device=batch.device,
-            assign_draft_cache_locs_page_size_1[(bs,)](
+                )
-                batch.req_pool_indices,
+                # FIXME(lsyin): align with the default code path
-                req_to_token_pool.req_to_token,
+                assign_draft_cache_locs_page_size_1[(bs,)](
-                batch.seq_lens,
+                    batch.req_pool_indices,
-                batch.out_cache_loc,
+                    req_to_token_pool.req_to_token,
-                req_to_token_pool.req_to_token.shape[1],
+                    batch.seq_lens,
-                topk,
+                    batch.out_cache_loc,
-                num_steps,
+                    req_to_token_pool.req_to_token.shape[1],
-            )
+                    topk,
+                    num_steps,
+                )
+            else:
+                # page_size > 1 + topk > 1: per-branch page-aligned draft pages.
+                # Reduce out_cache_loc from the page-aligned tree region down to the
+                # dense draft slots (skip each branch's duplicated prefix-tail slots
+                # and trailing padding), matching generate_draft_decode_kv_indices'
+                # paged read formula: prefix_base + t*num_new_pages*page + last_page + s.
+                # base is batch.seq_lens (== KV-ready committed prefix at draft time;
+                # the bonus is the tree root written by verify, not part of [0:seq_lens]).
+                rows = req_to_token_pool.req_to_token[batch.req_pool_indices.long()]
+                seq_lens = batch.seq_lens.to(torch.int64)
+                last_page = seq_lens % page_size
+                prefix_base = seq_lens - last_page
+                num_new_pages = (last_page + num_steps + page_size - 1) // page_size
+                topk_ids = torch.arange(
+                    topk, device=rows.device, dtype=torch.int64
+                ).view(1, topk)
+                starts = (
+                    prefix_base.view(bs, 1)
+                    + topk_ids * (num_new_pages.view(bs, 1) * page_size)
+                    + last_page.view(bs, 1)
+                )
+                steps = torch.arange(
+                    num_steps, device=rows.device, dtype=torch.int64
+                ).view(1, 1, num_steps)
+                pos = (starts.view(bs, topk, 1) + steps).reshape(bs, topk * num_steps)
+                batch.out_cache_loc = (
+                    torch.gather(rows, 1, pos).reshape(-1).contiguous()
+                )
+
+                # Each branch's page-aligned region starts with `last_page` hole slots
+                # overlapping the prefix tail page; duplicate the real prefix-tail KV
+                # into them so whole-page reads stay coherent (see helper docstring).
+                duplicate_prefix_tail_to_draft_branches(
+                    draft_model_runner.token_to_kv_pool,
+                    rows,
+                    prefix_base,
+                    last_page,
+                    num_new_pages,
+                    topk,
+                    page_size,
+                )
 
         # Get a forward batch
         self.num_tokens_per_req = topk

test/registered/spec/eagle/test_spec_eagle_page.py

@@ -1,7 +1,7 @@
-"""page_size > 1 variants (flashinfer).
+"""page_size > 1 variants at topk=1 (flashinfer).
 
-EAGLE3 page64 (spec v2) + EAGLE/Llama-2 page4 (topk1 and topk8, spec v1).
+EAGLE3 page64 (spec v2) + EAGLE/Llama-2 page4 (spec v1). topk>1 page variants
-Runs on the cheap (5090) runner.
+live in test_spec_eagle_topk.py. Runs on the cheap (5090) runner.
 """
 
 import unittest
@@ -15,7 +15,7 @@ from sglang.test.kits.spec_server_kits import (
 )
 from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
 
-register_cuda_ci(est_time=540, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=360, stage="base-b", runner_config="1-gpu-small")
 
 
 class TestEagle3Page64(Eagle3Base, SpecAccuracyKit, SpecLogprobKit, SpecFeatureKit):
@@ -35,12 +35,5 @@ class TestEagleLlama2Page4Topk1(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit
     env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
 
 
-class TestEagleLlama2Page4Topk8(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
-    """Llama-2 topk>1 tree + page_size=4 (spec v1)."""
-
-    page_size = 4
-    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
-
-
 if __name__ == "__main__":
     unittest.main()

test/registered/spec/eagle/test_spec_eagle_topk_page.py

@@ -0,0 +1,41 @@
+"""topk > 1 tree drafting at page_size > 1 (EAGLE3 topk8 + EAGLE/Llama-2 topk8).
+
+page64 stays on spec v2 (overlap), page4 runs on spec v1 (no overlap). flashinfer is
+pinned because this runs on the cheap (5090) runner, where fa3 (Hopper-only) isn't
+available -- functional sanity only, no perf/stress. (page>1 topk>1 on fa3 is covered
+on the Hopper runner in test_spec_eagle_fa3.py.)
+"""
+
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.spec_server_kits import (
+    SpecAccuracyKit,
+    SpecFeatureKit,
+)
+from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
+
+register_cuda_ci(est_time=720, stage="base-b", runner_config="1-gpu-small")
+
+
+class TestEagle3Page64Topk8(Eagle3Base, SpecAccuracyKit, SpecFeatureKit):
+    """EAGLE3 topk=8 tree + page_size=64 (spec v2)."""
+
+    page_size = 64
+    spec_topk = 8
+    spec_tokens = 32
+    disable_overlap = False
+    cuda_graph_max_bs = 5
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+class TestEagleLlama2Page4Topk8(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
+    """Llama-2 topk>1 tree + page_size=4 (spec v1)."""
+
+    page_size = 4
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+if __name__ == "__main__":
+    unittest.main()