[Perf][Spec Decoding] Skip cat/topk/sort/gather in draft_forward for topk=1 (#26424)

python/sglang/srt/speculative/eagle_utils.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, List, Optional
 import torch
 
 from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
+from sglang.srt.utils.async_probe import maybe_detect_oob
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -98,13 +99,21 @@ def organize_draft_results(
     top_scores = torch.topk(score_list, num_draft_token - 1, dim=-1)
     top_scores_index = top_scores.indices
     top_scores_index = torch.sort(top_scores_index).values
+    maybe_detect_oob(
+        top_scores_index,
+        0,
+        ss_token_list.shape[1],
+        "organize_draft_results: top_scores_index OOB for gather on ss_token_list",
+    )
     draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)
 
     if len(parents_list) > 1:
         parent_list = torch.cat(parents_list[:-1], dim=1)
     else:
         batch_size = parents_list[0].shape[0]
-        parent_list = torch.empty(batch_size, 0, device=parents_list[0].device)
+        parent_list = torch.empty(
+            batch_size, 0, dtype=torch.long, device=parents_list[0].device
+        )
 
     return parent_list, top_scores_index, draft_tokens
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -59,6 +59,7 @@ from sglang.srt.speculative.eagle_utils import (
     TreeMaskMode,
     _eagle_prefill_tail_tokens,
     build_tree_kernel_efficient,
+    organize_draft_results,
     per_step_draft_out_cache_loc,
 )
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -141,6 +142,11 @@ class EagleDraftWorker(BaseDraftWorker):
             server_args.speculative_algorithm
         )
 
+        # Pre-allocated constants for the topk=1 chain fast path in draft_forward.
+        self._topk1_parents_prealloc = None
+        self._topk1_score_indices_prealloc = None
+        self._rebuild_topk1_chain_buffers()
+
         # Do not capture cuda graph in `TpModelWorker` init,
         # will capture later with init_cuda_graphs()
         backup_disable_cuda_graph = server_args.disable_cuda_graph
@@ -214,6 +220,35 @@ class EagleDraftWorker(BaseDraftWorker):
 
         self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
 
+    def _rebuild_topk1_chain_buffers(self) -> None:
+        # For topk=1 the draft tree degenerates to a chain, so parent_list and
+        # top_scores_index are runtime-invariant. Must be rebuilt after any
+        # change to speculative_num_steps / speculative_num_draft_tokens.
+        if self.topk != 1:
+            return
+        # _override_worker_state can set both directly, bypassing the hook that
+        # pins this relation; the fast path is only valid when it holds.
+        assert self.speculative_num_draft_tokens == self.speculative_num_steps + 1, (
+            "topk=1 requires speculative_num_draft_tokens == speculative_num_steps + 1, "
+            f"got {self.speculative_num_draft_tokens} and {self.speculative_num_steps}"
+        )
+        num_steps = self.speculative_num_steps
+        sa = self.server_args
+        max_bs = max(
+            sa.cuda_graph_max_bs or 0,
+            sa.max_running_requests or 0,
+            1,
+        )
+        # A single-step chain has no parent entries (slow path drops the last
+        # step). repeat (not expand): the kernel reads these as contiguous.
+        parent_width = num_steps if num_steps > 1 else 0
+        self._topk1_parents_prealloc = torch.arange(
+            -1, parent_width - 1, dtype=torch.long, device=self.device
+        ).repeat(max_bs, 1)
+        self._topk1_score_indices_prealloc = torch.arange(
+            num_steps, dtype=torch.long, device=self.device
+        ).repeat(max_bs, 1)
+
     def init_token_map(self):
         # Load hot token ids
         if self.speculative_algorithm.is_eagle3():
@@ -554,33 +589,25 @@ class EagleDraftWorker(BaseDraftWorker):
             forward_batch.positions.add_(1)
 
         # Organize the results
-        score_list = torch.cat(score_list, dim=1).flatten(
+        if (
-            1
+            self.topk == 1
-        )  # b, n, topk; n= 1 + (num_steps-1) * self.topk
+            and token_list[0].shape[0] <= self._topk1_parents_prealloc.shape[0]
-        ss_token_list = torch.cat(
+        ):
-            token_list, dim=1
+            # Chain topology: draft_tokens = concat of per-step tokens; the
-        )  # b, (self.topk + (num_steps-1) * self.topk)
+            # full-length topk/sort/gather over score_list collapses to an
-        top_scores = torch.topk(
+            # identity. parent_list and top_scores_index are runtime-invariant
-            score_list, self.speculative_num_draft_tokens - 1, dim=-1
+            # constants pre-allocated on the worker. Oversized batches (rare,
-        )
+            # would silently truncate the slice) fall through to the slow path.
-        top_scores_index = top_scores.indices
+            bs = token_list[0].shape[0]
-        top_scores_index = torch.sort(top_scores_index).values
+            draft_tokens = torch.cat(token_list, dim=1)
-        maybe_detect_oob(
+            top_scores_index = self._topk1_score_indices_prealloc[:bs]
-            top_scores_index,
+            parent_list = self._topk1_parents_prealloc[:bs]
-            0,
-            ss_token_list.shape[1],
-            "draft_forward: top_scores_index OOB for gather on ss_token_list",
-        )
-        draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)
-
-        if len(parents_list) > 1:
-            parent_list = torch.cat(parents_list[:-1], dim=1)
-        else:
-            batch_size = parents_list[0].shape[0]
-            parent_list = torch.empty(batch_size, 0, device=parents_list[0].device)
-
             return parent_list, top_scores_index, draft_tokens
 
+        return organize_draft_results(
+            score_list, token_list, parents_list, self.speculative_num_draft_tokens
+        )
+
     def draft_extend(self):
         pass
 
@@ -1032,6 +1059,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
         dw.cuda_graph_runner = state.cuda_graph_runner
         dw.draft_extend_attn_backend = state.draft_extend_attn_backend
         dw.cuda_graph_runner_for_draft_extend = state.cuda_graph_runner_for_draft_extend
+        dw._rebuild_topk1_chain_buffers()
 
         # Target side
         self._target_worker.model_runner.attn_backend = state.target_attn_backend
@@ -1070,6 +1098,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
         dw.speculative_num_draft_tokens = speculative_num_draft_tokens
         sa.speculative_num_steps = speculative_num_steps
         sa.speculative_num_draft_tokens = speculative_num_draft_tokens
+        dw._rebuild_topk1_chain_buffers()
 
         try:
             yield
@@ -1087,6 +1116,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 sa.speculative_num_steps,
                 sa.speculative_num_draft_tokens,
             ) = backup
+            dw._rebuild_topk1_chain_buffers()
 
     def verify(self, batch: ScheduleBatch):
         fwd_stream = torch.get_device_module(self.device).current_stream()

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -70,6 +70,11 @@ class StandaloneDraftWorker(EagleDraftWorker):
             server_args.speculative_algorithm
         )
 
+        # Pre-allocated constants for the topk=1 chain fast path in draft_forward.
+        self._topk1_parents_prealloc = None
+        self._topk1_score_indices_prealloc = None
+        self._rebuild_topk1_chain_buffers()
+
         # Set constant
         from sglang.srt.speculative.eagle_info import EagleDraftInput
 

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -0,0 +1,94 @@
+"""Equivalence tests for the EagleDraftWorker topk=1 chain fast path.
+
+For topk=1 the draft tree degenerates to a chain, so `draft_forward` skips the
+cat/topk/sort/gather of the slow path and returns pre-allocated constants. These
+tests check that the pre-allocated `parent_list` / `top_scores_index` match the
+slow path (`organize_draft_results`) for num_steps in {1, 2, 3, 4}.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.speculative.eagle_utils import organize_draft_results
+from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+from sglang.srt.utils import get_device
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, suite="base-b-test-1-gpu-small")
+
+DEVICE = get_device()
+
+
+def _make_chain_lists(num_steps: int, bs: int):
+    """Build the (score, token, parents) lists a topk=1 chain produces.
+
+    Shapes/values mirror `select_top_k_tokens` for topk=1: each step yields one
+    token; the first step's parents are [-1, 0], later steps' parents are [i].
+    """
+    score_list, token_list, parents_list = [], [], []
+    for i in range(num_steps):
+        # Strictly decreasing scores, as a real chain produces (cumulative probs).
+        score_list.append(torch.full((bs, 1, 1), float(num_steps - i), device=DEVICE))
+        token_list.append(
+            torch.arange(i * bs, (i + 1) * bs, device=DEVICE).unsqueeze(1)
+        )
+        if i == 0:
+            parents_list.append(
+                torch.tensor([-1, 0], dtype=torch.long, device=DEVICE).repeat(bs, 1)
+            )
+        else:
+            parents_list.append(torch.full((bs, 1), i, dtype=torch.long, device=DEVICE))
+    return score_list, token_list, parents_list
+
+
+def _make_worker(num_steps: int, num_draft_tokens: int):
+    worker = object.__new__(EagleDraftWorker)
+    worker.topk = 1
+    worker.device = DEVICE
+    worker.speculative_num_steps = num_steps
+    worker.speculative_num_draft_tokens = num_draft_tokens
+    worker.server_args = SimpleNamespace(cuda_graph_max_bs=8, max_running_requests=8)
+    return worker
+
+
+class TestEagleWorkerV2Topk1FastPath(CustomTestCase):
+    def test_fast_path_matches_slow_path(self):
+        bs = 3
+        for num_steps in (1, 2, 3, 4):
+            with self.subTest(num_steps=num_steps):
+                num_draft_tokens = num_steps + 1
+                worker = _make_worker(num_steps, num_draft_tokens)
+                worker._rebuild_topk1_chain_buffers()
+
+                score_list, token_list, parents_list = _make_chain_lists(num_steps, bs)
+                ref_parent, ref_index, ref_tokens = organize_draft_results(
+                    score_list, token_list, parents_list, num_draft_tokens
+                )
+
+                fast_parent = worker._topk1_parents_prealloc[:bs]
+                fast_index = worker._topk1_score_indices_prealloc[:bs]
+                fast_tokens = torch.cat(token_list, dim=1)
+
+                self.assertEqual(fast_parent.shape, ref_parent.shape)
+                self.assertEqual(fast_parent.tolist(), ref_parent.long().tolist())
+                self.assertEqual(fast_index.tolist(), ref_index.long().tolist())
+                self.assertEqual(fast_tokens.tolist(), ref_tokens.tolist())
+
+                # The kernel reads these via data_ptr() as contiguous int64.
+                self.assertEqual(fast_parent.dtype, torch.long)
+                self.assertEqual(fast_index.dtype, torch.long)
+                self.assertTrue(fast_parent.is_contiguous())
+                self.assertTrue(fast_index.is_contiguous())
+
+    def test_assert_on_inconsistent_steps_and_draft_tokens(self):
+        # num_draft_tokens must equal num_steps + 1 for topk=1.
+        worker = _make_worker(num_steps=3, num_draft_tokens=3)
+        with self.assertRaises(AssertionError):
+            worker._rebuild_topk1_chain_buffers()
+
+
+if __name__ == "__main__":
+    unittest.main()