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

[NPU] Enable non-greedy MTP sampling (#32495)

Browse Source
This commit is contained in:
Wuhen Duan
2026-09-08 11:18:01 +08:00
committed by GitHub
parent 28ebede865
commit dfd9b5c2a4
3 changed files with 37 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_utils.py
+17 -9
View File
@@ -781,7 +781,7 @@ def eagle_sample(
# Sample tokens
target_predict = None
if sampling_info.is_all_greedy or _is_cpu or _is_npu or _is_hip or _is_xpu:
if sampling_info.is_all_greedy or _is_cpu or _is_hip or _is_xpu:
target_predict = torch.argmax(next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, verify_input.draft_token_num)
predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
@@ -847,15 +847,23 @@ def eagle_sample(
tp_group.broadcast(accept_index, src=0)
tp_group.broadcast(num_correct_drafts, src=0)
else:
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
if _is_npu:
from sgl_kernel_npu.sample import (
chain_speculative_sampling_triton,
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
else:
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
from sglang.kernels.ops.speculative.reject_sampling import (
chain_speculative_sampling_triton,
)
from sglang.kernels.ops.speculative.reject_sampling import (
chain_speculative_sampling_triton,
)
use_rejection_sampling = get_spec().speculative_use_rejection_sampling