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

Skip no-op EAGLE sampling renormalization (#31294)

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-07-15 22:37:10 -07:00
committed by GitHub
co-authored by weireweire
parent 40517b593b
commit ac4fa6496c
1 changed files with 16 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_utils.py
+16 -14
View File
@@ -683,20 +683,22 @@ def eagle_sample(
next_token_logits / expanded_temperature, dim=-1
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
target_probs = top_k_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ks, verify_input.draft_token_num, dim=0
),
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
target_probs = top_p_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ps, verify_input.draft_token_num, dim=0
),
)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm")
if sampling_info.need_top_k_sampling:
target_probs = top_k_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ks, verify_input.draft_token_num, dim=0
),
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
if sampling_info.need_top_p_sampling:
target_probs = top_p_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ps, verify_input.draft_token_num, dim=0
),
)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm")
target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1)
draft_probs = (
verify_input.draft_probs