From e99729c9f30fdfa9803421e9792f650f8de3925a Mon Sep 17 00:00:00 2001 From: jiapingW <56055330+jiapingW@users.noreply.github.com> Date: Thu, 21 Aug 2025 07:42:01 +0800 Subject: [PATCH] Fixed the issue where eagle3 TPOT was not as good as without eagle3. (#9404) --- python/sglang/srt/speculative/eagle_utils.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index d4741144d..b02319584 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -453,12 +453,13 @@ class EagleVerifyInput: sampling_info.top_ks, self.draft_token_num, dim=0 ), ) # (bs * draft_token_num, vocab_size) - target_probs = top_p_renorm_prob( - target_probs, - torch.repeat_interleave( - sampling_info.top_ps, self.draft_token_num, dim=0 - ), - ) + if not torch.all(sampling_info.top_ps == 1.0): + target_probs = top_p_renorm_prob( + target_probs, + torch.repeat_interleave( + sampling_info.top_ps, self.draft_token_num, dim=0 + ), + ) target_probs = target_probs.reshape(bs, self.draft_token_num, -1) draft_probs = torch.zeros(