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
docs/docs/hardware-platforms/ascend-npus/optimization/parameter_tuning.mdx
+2
View File
@@ -767,6 +767,8 @@ Captures the compute graph on NPU and replays it to eliminate kernel launch over
Reduces per-token latency by predicting multiple future tokens in a single forward pass, then verifying them against the model. Ascend supports `NEXTN` (for DeepSeek models, using the model's own hidden states) and `EAGLE3` (for Qwen MoE models, using a separate draft model). Controlled by `--speculative-num-steps`, `--speculative-eagle-topk`, and `--speculative-num-draft-tokens`. On Ascend, `SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1` enables the optimized overlap scheduler for speculative decoding.
For non-greedy requests, Ascend target verification applies request-level `temperature`, `top_k`, and `top_p`. The default target-only mode supports linear and tree-shaped candidates. Classic rejection sampling is available with `--speculative-use-rejection-sampling` for linear chains configured with `--speculative-eagle-topk 1`.
### PrefixCache (`--disable-radix-cache`)
Reuses KV cache across requests that share common prompt prefixes (Radix Cache), reducing repeated prefill computation and lowering time-to-first-token. Enabled by default; disable with `--disable-radix-cache` when prefix reuse is not expected (e.g., dedicated prefill nodes in PD disaggregation, or random-input benchmarks).
docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx
+18 -2
View File
@@ -1373,6 +1373,16 @@ click [Server Arguments](../../../advanced_features/server_arguments).
## Speculative decoding
Ascend NPU supports non-greedy target verification for the supported EAGLE-based speculative algorithms. Request-level
`temperature`, `top_k`, and `top_p` are applied when constructing the target probability distribution. Request-level
`top_k` controls token sampling and is separate from `--speculative-eagle-topk`, which controls the draft tree branching
factor.
The default target-only mode supports both linear and tree-shaped draft candidates. Classic rejection sampling is
available for EAGLE3 and NEXTN with `--speculative-use-rejection-sampling` and requires
`--speculative-eagle-topk 1`. Rejection sampling uses the draft proposal distribution and cannot be combined with
non-default speculative acceptance thresholds or deterministic inference.
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "16.7%"}} />
@@ -1443,13 +1453,19 @@ click [Server Arguments](../../../advanced_features/server_arguments).
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--speculative-accept-threshold-single`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`1.0`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Type: float</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--speculative-accept-threshold-acc`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`1.0`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Type: float</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--speculative-use-rejection-sampling`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag<br/> (requires `--speculative-eagle-topk 1`)</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--speculative-token-map`</td>
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