chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Tiny remove the duplicate function in spec v2 (#14957)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-12-13 12:08:27 +09:00
committed by GitHub
parent 77873343c4
commit 7160283800
2 changed files with 3 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -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,

4
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -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])