Enable piecewise-cuda-graph when logprob_start_len = -1 (#19453)

This commit is contained in:
Qiaolin Yu
2026-03-10 12:50:57 -07:00
committed by GitHub
parent 031d0a2aad
commit a3d88a247b
3 changed files with 4 additions and 4 deletions

View File

@@ -1157,7 +1157,7 @@ class Req(ReqDllmMixin):
# - extend_input_len: Number of tokens that need to be processed in this extend batch
self.extend_input_len = extend_input_len
if self.logprob_start_len == -1:
logprob_start_len = len(self.fill_ids) - 1
logprob_start_len = len(self.fill_ids)
else:
# logprob_start_len should be at least the length of the prefix indices
logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices))
@@ -1582,7 +1582,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
len(req.fill_ids),
)
if req.logprob_start_len == -1:
logprob_start_len = len(req.origin_input_ids) - 1
logprob_start_len = len(req.origin_input_ids)
else:
logprob_start_len = req.logprob_start_len
# Apply logprob_start_len

View File

@@ -1776,7 +1776,7 @@ class Scheduler(
if recv_req.return_logprob and recv_req.token_ids_logprob is None:
# If logprob is required but neither token_ids_logprob nor logprob_start_len is
# set, return the logprobs for output tokens by default
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = len(req.origin_input_ids)
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