Skip no-op EAGLE sampling renormalization (#31294)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
2026-07-15 22:37:10 -07:00
parent 40517b593b
commit ac4fa6496c
1 changed files with 16 additions and 14 deletions
@@ -683,6 +683,7 @@ 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")
+        if sampling_info.need_top_k_sampling:
            target_probs = top_k_renorm_prob(
                target_probs,
                torch.repeat_interleave(
@@ -690,6 +691,7 @@ def eagle_sample(
                ),
            )  # (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(