Support spec v2 tree drafting (eagle topk>1) with page_size==1 (#26866)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -278,10 +278,14 @@ 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
         and not server_args.disable_overlap_schedule
     ):
+        # Spec v2 tree drafting supports topk > 1 with page_size == 1. The
+        # page_size > 1 + topk > 1 draft KV allocation (partial-page duplication)
+        # is not yet ported to v2, so fall back to v1 only for that case.
         server_args.disable_overlap_schedule = True
-        spec_v1_reason = "spec v2 currently only supports topk = 1"
+        spec_v1_reason = "spec v2 topk > 1 currently requires page_size == 1"
     elif (
         not envs.SGLANG_ENABLE_SPEC_V2.get()
         and not server_args.disable_overlap_schedule

python/sglang/srt/speculative/eagle_info_v2.py

@@ -506,7 +506,7 @@ def fill_bonus_tokens(
     accept_tokens,
     accept_lens,
     bonus_tokens_ptr,
-    num_draft_tokens: tl.constexpr,
+    accept_stride: tl.constexpr,
 ):
     # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel
     # because this kernel reads accept_lens
@@ -514,7 +514,8 @@ def fill_bonus_tokens(
     # `accept_lens` includes the bonus token; the last accepted slot is at -1.
     accept_len = tl.load(accept_lens + pid)
 
-    bonus_token_idx = num_draft_tokens * pid + accept_len - 1
+    # accept_stride = per-req width of accept_tokens (= accept_index.shape[1]).
+    bonus_token_idx = accept_stride * pid + accept_len - 1
     bonus_token = tl.load(accept_tokens + bonus_token_idx)
     tl.store(bonus_tokens_ptr + pid, bonus_token)
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -1202,11 +1202,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
         if not batch.forward_mode.is_idle():
             accept_tokens = predict[accept_index]
             bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32)
+            # stride = accept_tokens per-req width = accept_index.shape[1]
+            # (spec_steps + 1); NOT num_draft_tokens, wrong for topk > 1 trees.
             fill_bonus_tokens[(bs,)](
                 accept_tokens,
                 accept_lens,
                 bonus_tokens,
-                self.speculative_num_draft_tokens,
+                accept_index.shape[1],
             )
         else:
             bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)
@@ -1216,6 +1218,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 batch, logits_output, predict, accept_index, self.speculative_num_steps
             )
 
+        if not batch.forward_mode.is_idle() and self.topk > 1:
+            # topk == 1 needs nothing here: the accepted path is already the front
+            # chain, so the whole compaction is an identity transform.
+            predict = self._finalize_accepted_tree_path(
+                batch, accept_index, accept_lens, predict, logits_output, bs
+            )
+
         next_draft_input = EagleDraftInput(bonus_tokens=bonus_tokens)
 
         # verify_forward_batch transitively holds verify-time GPU tensors
@@ -1297,6 +1306,30 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 model=self.target_worker.model_runner.model,
             )
 
+    def _finalize_accepted_tree_path(
+        self,
+        batch: ScheduleBatch,
+        accept_index: torch.Tensor,
+        accept_lens: torch.Tensor,
+        predict: torch.Tensor,
+        logits_output,
+        bs: int,
+    ) -> torch.Tensor:
+        """Tree drafting (topk > 1): move the accepted path -- KV slots, predict,
+        hidden_states -- to the contiguous front of each per-req block, which the
+        downstream chain-layout code (draft-extend select_index, committed-KV reads)
+        assumes. Returns compacted predict; mutates logits_output.hidden_states
+        (moved only when present)."""
+        self.move_accepted_tokens_to_target_kvcache(
+            batch, accept_index, accept_lens - 1
+        )
+        predict = self._compact_accepted_to_front(predict, accept_index, bs)
+        if logits_output.hidden_states is not None:
+            logits_output.hidden_states = self._compact_accepted_to_front(
+                logits_output.hidden_states, accept_index, bs
+            )
+        return predict
+
     def move_accepted_tokens_to_target_kvcache(
         self,
         batch: ScheduleBatch,
@@ -1313,7 +1346,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot.
         """
         bs = len(batch.seq_lens)
-        size = bs * self.speculative_num_draft_tokens
+        # accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the
+        # tree exceeds the accepted chain, over-reading accept_index (illegal memory).
+        size = bs * accept_index.shape[1]
 
         # fill_accepted_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
         maybe_detect_oob(
@@ -1350,6 +1385,24 @@ class EAGLEWorkerV2(BaseSpecWorker):
             tgt_cache_loc, accepted_out_cache_loc
         )
 
+    def _compact_accepted_to_front(
+        self, x: torch.Tensor, accept_index: torch.Tensor, bs: int
+    ) -> torch.Tensor:
+        """Gather the accepted tree path to the front of each per-req block.
+
+        ``x`` is node-indexed over the whole tree (``[bs * num_draft_tokens, ...]``),
+        ``accept_index`` is ``[bs, spec_steps + 1]`` global node indices (-1 padded).
+        Padded entries clamp to node 0 but land past accept_lens (never read);
+        trailing unaccepted slots stay and are freed as overshoot.
+        """
+        nd = self.speculative_num_draft_tokens
+        s1 = accept_index.shape[1]  # spec_steps + 1
+        safe = accept_index.to(torch.int64).clamp(min=0).reshape(-1)
+        gathered = x[safe]
+        out = x.clone()
+        out.view(bs, nd, *x.shape[1:])[:, :s1] = gathered.view(bs, s1, *x.shape[1:])
+        return out
+
     def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
         success, message = self._draft_worker.draft_runner.update_weights_from_disk(
             recv_req.model_path,

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -790,11 +790,12 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
         if not batch.forward_mode.is_idle():
             accept_tokens = predict[accept_index]
             bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32)
+            # stride = accept_tokens per-req width = accept_index.shape[1].
             fill_bonus_tokens[(bs,)](
                 accept_tokens,
                 accept_lens,
                 bonus_tokens,
-                self.speculative_num_draft_tokens,
+                accept_index.shape[1],
             )
         else:
             bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)

test/registered/spec/eagle/test_spec_eagle_stress.py

@@ -20,7 +20,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=600, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=780, stage="base-b", runner_config="1-gpu-large")
 
 
 class TestEagle3Perf(Eagle3Base, SpecPerfKit):
@@ -40,6 +40,20 @@ class TestEagleLlama2Retract(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
     )
 
 
+class TestEagle3Topk16V2Retract(Eagle3Base, SpecAccuracyKit, SpecFeatureKit):
+    """EAGLE3 topk=16 tree on spec v2 under retract; must not leak KV. Stresses
+    the accepted-path KV move (move_accepted_tokens_to_target_kvcache)."""
+
+    spec_topk = 16
+    spec_tokens = 64
+    disable_overlap = False
+    cuda_graph_max_bs = 5
+    max_running_requests = 64
+    gsm8k_accept_len_thres = 2.4
+    extra_args = ("--max-total-tokens", 4500)  # small KV to trigger retract
+    env_overrides = ((envs.SGLANG_TEST_RETRACT, True),)
+
+
 class TestEagleLlama2AbortAll(EagleLlama2Base, AbortAllMixin):
     abort_all_max_new_tokens = 4000
 

test/registered/spec/eagle/test_spec_eagle_topk.py

@@ -1,7 +1,9 @@
 """topk > 1 tree drafting (EAGLE3 topk16 + EAGLE/Llama-2 topk8).
 
-topk > 1 always routes to spec v1; flashinfer is pinned (topk > 1 can't use fa3).
+topk > 1 routes to spec v1, except page_size==1 which can also stay on spec v2
-Runs on the cheap (5090) runner -- functional sanity only, no perf/stress.
+(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.
+(topk > 1 on fa3 is covered on the Hopper runner in test_spec_eagle_fa3.py.)
 """
 
 import unittest
@@ -16,7 +18,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=840, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=1180, stage="base-b", runner_config="1-gpu-small")
 
 
 class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogprobKit):
@@ -31,6 +33,13 @@ class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogp
     gsm8k_accept_len_thres = 2.4  # EAGLE3 topk16 gsm8k accept ~2.48
 
 
+class TestEagle3Topk16SpecV2(TestEagle3Topk16, SpecFeatureKit):
+    """EAGLE3 topk=16 tree on spec v2 (overlap, page1): guards the v2 tree path's
+    accepted-path compaction, validated by logprob_spec_v2_match."""
+
+    disable_overlap = False
+
+
 class TestEagleLlama2Suite(
     EagleLlama2Base,
     SpecCorrectnessKit,