From 9f5655340897ae0af9425acd528e809cf223c2dc Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 30 Jul 2026 02:55:21 -0700 Subject: [PATCH] [Perf] Fast-path chain-style draft token organization in multi-layer EAGLE (#32887) --- .../srt/speculative/base_spec_worker.py | 40 ++++++++++++++ python/sglang/srt/speculative/eagle_utils.py | 2 + .../sglang/srt/speculative/eagle_worker_v2.py | 37 ------------- .../multi_layer_eagle_worker_v2.py | 55 ++++++++++--------- .../srt/speculative/standalone_worker_v2.py | 3 - 5 files changed, 71 insertions(+), 66 deletions(-) diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index d2bddc14e..eb2119413 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -3,6 +3,8 @@ from __future__ import annotations from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Optional +import torch + if TYPE_CHECKING: from sglang.srt.managers.io_struct import ( UpdateWeightFromDiskReqInput, @@ -13,6 +15,10 @@ if TYPE_CHECKING: class EagleDraftWorkerBase(ABC): + # topk=1 chain constants for draft_forward's fast path; None when topk > 1. + _topk1_parents_prealloc: Optional[torch.Tensor] = None + _topk1_score_indices_prealloc: Optional[torch.Tensor] = None + @abstractmethod def draft(): pass @@ -41,6 +47,40 @@ class EagleDraftWorkerBase(ABC): self.draft_worker.init_cuda_graphs(capture_decode_cuda_graph=False) self._capture_cuda_graphs() + 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 + decode_max_bs = ( + sa.cuda_graph_config.decode.max_bs + if sa.cuda_graph_config is not None + else None + ) + max_bs = max( + decode_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) + class BaseSpecWorker(ABC): @property diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 6e25dd8b1..c9037aaf1 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -103,7 +103,9 @@ def organize_draft_results( parents_list: List[torch.Tensor], num_draft_token: int, ): + # b, n, topk; n = 1 + (num_steps-1) * topk score_list = torch.cat(score_list, dim=1).flatten(1) + # b, (topk + (num_steps-1) * topk) ss_token_list = torch.cat(token_list, dim=1) top_scores = torch.topk(score_list, num_draft_token - 1, dim=-1) top_scores_index = top_scores.indices diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 3c743f7fb..ee4357578 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -144,9 +144,6 @@ class EagleDraftWorker(EagleDraftWorkerBase): 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() # Load draft model weights only. @@ -251,40 +248,6 @@ class EagleDraftWorker(EagleDraftWorkerBase): self.index_share_for_mtp_iteration and self.dsa_index_topk is not None ) - 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 - decode_max_bs = ( - sa.cuda_graph_config.decode.max_bs - if sa.cuda_graph_config is not None - else None - ) - max_bs = max( - decode_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(): diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 0bc2620e2..2b2e15578 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -50,6 +50,7 @@ from sglang.srt.speculative.eagle_info import ( from sglang.srt.speculative.eagle_utils import ( default_tree_mask_mode, get_draft_recurrent_hidden_state_spec, + organize_draft_results, ) from sglang.srt.speculative.eagle_worker_common import ( build_eagle_verify_input, @@ -133,6 +134,8 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): server_args.speculative_algorithm ) + self._rebuild_topk1_chain_buffers() + # Set constant EagleDraftInput.ALLOC_LEN_PER_DECODE = max( self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens @@ -442,6 +445,30 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info") + # Chain-style (topk=1, one token per draft step, all of them selected): + # _draft_forward_organize's slice/cat/topk/sort/gather is the identity on + # topk_index, and parent_list is the constant [-1, 0, .., S-2] per row. + parents_prealloc = self._topk1_parents_prealloc + if ( + parents_prealloc is not None + and topk_index.shape[1] == self.speculative_num_steps + and topk_index.shape[0] <= parents_prealloc.shape[0] + ): + bs = topk_index.shape[0] + return ( + parents_prealloc[:bs], + self._topk1_score_indices_prealloc[:bs], + topk_index, + ) + + return self._draft_forward_organize(topk_p, topk_index, hidden_states) + + def _draft_forward_organize( + self, + topk_p: torch.Tensor, + topk_index: torch.Tensor, + hidden_states: torch.Tensor, + ): # Return values score_list: List[torch.Tensor] = [] token_list: List[torch.Tensor] = [] @@ -473,33 +500,9 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): ) ) - # 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 + return organize_draft_results( + score_list, token_list, parents_list, self.speculative_num_draft_tokens ) - 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 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 def draft_extend(self): pass diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 624fb5169..ce04a28ef 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -51,9 +51,6 @@ 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