Skip no-op EAGLE sampling renormalization (#31294)

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
@@ -683,6 +683,7 @@ def eagle_sample(
next_token_logits / expanded_temperature, dim=-1 next_token_logits / expanded_temperature, dim=-1
) # (bs * num_draft_tokens, vocab_size) ) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
if sampling_info.need_top_k_sampling:
target_probs = top_k_renorm_prob( target_probs = top_k_renorm_prob(
target_probs, target_probs,
torch.repeat_interleave( torch.repeat_interleave(
@@ -690,6 +691,7 @@ def eagle_sample(
), ),
) # (bs * num_draft_tokens, vocab_size) ) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") 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 = top_p_renorm_prob(
target_probs, target_probs,
torch.repeat_interleave( torch.repeat_interleave(