[PD] Fix top logprobs crash in prefill path (#26299)

This commit is contained in:
Jun Liu
2026-05-26 22:01:10 +08:00
committed by GitHub
parent c8c1aed5e9
commit c47f0e7cdd
2 changed files with 37 additions and 0 deletions
@@ -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)