Revert "Support spec v2 tree drafting (eagle topk>1) with page_size==1" (#26981)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -278,14 +278,10 @@ 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 topk > 1 currently requires page_size == 1"
+        spec_v1_reason = "spec v2 currently only supports topk = 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,
-    accept_stride: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
 ):
     # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel
     # because this kernel reads accept_lens
@@ -514,8 +514,7 @@ def fill_bonus_tokens(
     # `accept_lens` includes the bonus token; the last accepted slot is at -1.
     accept_len = tl.load(accept_lens + pid)
 
-    # accept_stride = per-req width of accept_tokens (= accept_index.shape[1]).
-    bonus_token_idx = accept_stride * pid + accept_len - 1
+    bonus_token_idx = num_draft_tokens * 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

@@ -1232,13 +1232,11 @@ 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,
-                accept_index.shape[1],
+                self.speculative_num_draft_tokens,
             )
         else:
             bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)
@@ -1248,13 +1246,6 @@ 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
@@ -1336,30 +1327,6 @@ 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,
@@ -1376,9 +1343,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot.
         """
         bs = len(batch.seq_lens)
-        # 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]
+        size = bs * self.speculative_num_draft_tokens
 
         # fill_accepted_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
         maybe_detect_oob(
@@ -1415,24 +1380,6 @@ 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,12 +790,11 @@ 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,
-                accept_index.shape[1],
+                self.speculative_num_draft_tokens,
             )
         else:
             bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)