Fix logprob_start_len handling for prefill-only requests (#17395)

This commit is contained in:
Cheng Wan
2026-01-29 15:14:43 -08:00
committed by GitHub
parent c3bf53c7c1
commit 6a6b36367e

View File

@@ -1545,13 +1545,13 @@ class Scheduler(
recv_req.logprob_start_len = -1
if recv_req.logprob_start_len == -1:
if req.is_prefill_only:
if recv_req.return_logprob:
# If return_logprob is True, return the logprobs for output tokens by default
req.logprob_start_len = len(req.origin_input_ids) - 1
elif req.is_prefill_only:
# For prefill-only requests with logprob_start_len == -1, set logprob_start_len
# beyond input sequence to skip input logprob computation entirely
req.logprob_start_len = len(req.origin_input_ids)
elif recv_req.return_logprob:
# If return_logprob is True, return the logprobs for output tokens by default
req.logprob_start_len = len(req.origin_input_ids) - 1
else:
# If return_logprob is False, only the last token requires logprob computation
req.logprob_start_len = -1