Fix overlap scheduler not take effect when outputing logprobs (#14096)
This commit is contained in:
@@ -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 = (
|
||||
|
||||
Reference in New Issue
Block a user