[PD] Fix top logprobs crash in prefill path (#26299)
This commit is contained in:
@@ -510,6 +510,17 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
logits_output.input_token_logprobs = tuple(
|
||||
logits_output.input_token_logprobs.tolist()
|
||||
)
|
||||
if logits_output.next_token_top_logprobs_val:
|
||||
logits_output.next_token_top_logprobs_val = [
|
||||
v.tolist() for v in logits_output.next_token_top_logprobs_val
|
||||
]
|
||||
logits_output.next_token_top_logprobs_idx = [
|
||||
x.tolist() for x in logits_output.next_token_top_logprobs_idx
|
||||
]
|
||||
if logits_output.next_token_token_ids_logprobs_val:
|
||||
logits_output.next_token_token_ids_logprobs_val = [
|
||||
v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
|
||||
]
|
||||
|
||||
for i, (req, next_token_id) in enumerate(
|
||||
zip(batch.reqs, next_token_ids, strict=True)
|
||||
|
||||
Reference in New Issue
Block a user