From 71602838007628f784c49eb51bc5543c9b05be28 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 13 Dec 2025 12:08:27 +0900 Subject: [PATCH] Tiny remove the duplicate function in spec v2 (#14957) --- .../sglang/srt/speculative/eagle_info_v2.py | 53 +------------------ .../sglang/srt/speculative/eagle_worker_v2.py | 4 +- 2 files changed, 3 insertions(+), 54 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 952ce740c..7c13ddc08 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -28,7 +28,7 @@ from sglang.srt.speculative.spec_utils import ( SIMULATE_ACC_LEN, generate_simulated_accept_index, ) -from sglang.srt.utils.common import fast_topk, is_cuda, is_hip, is_npu, next_power_of_2 +from sglang.srt.utils.common import is_cuda, is_hip, is_npu, next_power_of_2 _is_cuda = is_cuda() _is_hip = is_hip() @@ -47,7 +47,6 @@ if is_cuda(): top_p_renorm_prob, tree_speculative_sampling_target_only, ) - from sgl_kernel.top_k import fast_topk @triton.jit @@ -371,56 +370,6 @@ class EagleVerifyInputV2Mixin: return predict, accept_length, accept_index -@torch.compile(dynamic=True, disable=_is_npu) -def select_top_k_tokens_tmp( - i: int, - topk_p: torch.Tensor, - topk_index: torch.Tensor, - hidden_states: torch.Tensor, - scores: torch.Tensor, - topk: int, -): - # FIXME(lsyin): remove this duplicate code - if i == 0: - # The first step after extend - input_ids = topk_index.flatten() - hidden_states = hidden_states.repeat_interleave(topk, dim=0) - scores = topk_p # shape: (b, topk) - - tree_info = ( - topk_p.unsqueeze(1), # shape: (b, 1, topk) - topk_index, # shape: (b, topk) - torch.arange(-1, topk, dtype=torch.long, device=hidden_states.device) - .unsqueeze(0) - .repeat(topk_p.shape[0], 1), # shape: (b, topk + 1) - ) - else: - # The later decode steps - expand_scores = torch.mul( - scores.unsqueeze(2), topk_p.reshape(-1, topk, topk) - ) # (b, topk, 1) x (b, topk ,topk) -> (b, topk, topk) - topk_cs_p, topk_cs_index = fast_topk( - expand_scores.flatten(start_dim=1), topk, dim=-1 - ) # (b, topk) - scores = topk_cs_p # shape: (b, topk) - - topk_index = topk_index.reshape(-1, topk**2) - input_ids = torch.gather(topk_index, index=topk_cs_index, dim=1).flatten() - - selected_input_index = topk_cs_index.flatten() // topk + torch.arange( - 0, hidden_states.shape[0], step=topk, device=hidden_states.device - ).repeat_interleave(topk) - hidden_states = hidden_states[selected_input_index, :] - - tree_info = ( - expand_scores, # shape: (b, topk, topk) - topk_index, # shape: (b, topk * topk) - topk_cs_index + (topk**2 * (i - 1) + topk), # shape: (b, topk) - ) - - return input_ids, hidden_states, scores, tree_info - - @triton.jit def fill_new_verified_id( verified_id, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 629ddf27e..ce68e3cdd 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -35,7 +35,6 @@ from sglang.srt.speculative.eagle_info_v2 import ( assign_extend_cache_locs, fill_accepted_out_cache_loc, fill_new_verified_id, - select_top_k_tokens_tmp, ) from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -44,6 +43,7 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, generate_token_bitmask, load_token_map, + select_top_k_tokens, ) from sglang.srt.utils.common import ( MultiprocessingSerializer, @@ -372,7 +372,7 @@ class EagleDraftWorker(BaseDraftWorker): # Forward multiple steps scores = None for i in range(self.speculative_num_steps): - input_ids, hidden_states, scores, tree_info = select_top_k_tokens_tmp( + input_ids, hidden_states, scores, tree_info = select_top_k_tokens( i, topk_p, topk_index, hidden_states, scores, self.topk ) score_list.append(tree_info[0])