From 4151a04d1aadaae3473e62e010415cdb5e91bc83 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Mon, 1 Jun 2026 15:37:47 -0700 Subject: [PATCH] [Perf][Spec Decoding] Skip cat/topk/sort/gather in draft_forward for topk=1 (#26424) --- python/sglang/srt/speculative/eagle_utils.py | 11 ++- .../sglang/srt/speculative/eagle_worker_v2.py | 80 +++++++++++----- .../srt/speculative/standalone_worker_v2.py | 5 + .../test_eagle_worker_v2_topk1_fastpath.py | 94 +++++++++++++++++++ 4 files changed, 164 insertions(+), 26 deletions(-) create mode 100644 test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index db435075a..451a46dc9 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 6a841ed53..6e606ac37 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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,32 +589,24 @@ class EagleDraftWorker(BaseDraftWorker): forward_batch.positions.add_(1) # Organize the results - score_list = torch.cat(score_list, dim=1).flatten( - 1 - ) # b, n, topk; n= 1 + (num_steps-1) * self.topk - ss_token_list = torch.cat( - token_list, dim=1 - ) # b, (self.topk + (num_steps-1) * self.topk) - top_scores = torch.topk( - score_list, self.speculative_num_draft_tokens - 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], - "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 ( + self.topk == 1 + and token_list[0].shape[0] <= self._topk1_parents_prealloc.shape[0] + ): + # Chain topology: draft_tokens = concat of per-step tokens; the + # full-length topk/sort/gather over score_list collapses to an + # identity. parent_list and top_scores_index are runtime-invariant + # constants pre-allocated on the worker. Oversized batches (rare, + # would silently truncate the slice) fall through to the slow path. + bs = token_list[0].shape[0] + draft_tokens = torch.cat(token_list, dim=1) + top_scores_index = self._topk1_score_indices_prealloc[:bs] + parent_list = self._topk1_parents_prealloc[:bs] + return parent_list, top_scores_index, draft_tokens - 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() diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 5d39f2d97..84a59515d 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/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 diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py new file mode 100644 index 000000000..fe24a63f2 --- /dev/null +++ b/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()