Fix overlap scheduler not take effect when outputing logprobs (#14096)

This commit is contained in:
fzyzcjy
2025-11-28 18:15:56 +08:00
committed by GitHub
parent 0e8ce1e832
commit 45cf575852
2 changed files with 6 additions and 6 deletions

View File

@@ -43,15 +43,15 @@ class GenerationBatchResult:
# relay path: forward stream -> next step forward
next_draft_input: Optional[EagleDraftInput] = None
def copy_to_cpu(self, return_logprob: bool = False):
def copy_to_cpu(self, return_logprob: bool):
"""Copy tensors to CPU in overlap scheduling.
Only the tensors which are needed for processing results are copied,
e.g., next_token_ids, logits outputs
"""
if return_logprob:
if self.logits_output.next_token_logits is not None:
self.logits_output.next_token_logits = (
self.logits_output.next_token_logits.to("cpu", non_blocking=True)
if self.logits_output.next_token_logprobs is not None:
self.logits_output.next_token_logprobs = (
self.logits_output.next_token_logprobs.to("cpu", non_blocking=True)
)
if self.logits_output.input_token_logprobs is not None:
self.logits_output.input_token_logprobs = (