perf: avoid unnecessary gpu-cpu sync in eagle_info (#20266)

Co-authored-by: root <qianhao@zhejianglab.org>
2026-03-19 19:37:29 -07:00
parent 3d749c49ca
commit fa0d8f6629
1 changed files with 1 additions and 1 deletions
@@ -339,7 +339,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
                    sampling_info.top_ks, self.draft_token_num, dim=0
                ),
            )  # (bs * draft_token_num, vocab_size)
-            if not torch.all(sampling_info.top_ps == 1.0):
+            if sampling_info.need_top_p_sampling:
                target_probs = top_p_renorm_prob(
                    target_probs,
                    torch.repeat_interleave(