chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[scheduler] fix: correcting extend_logprob_start_len calculation (#15922)

Browse Source
This commit is contained in:
Cheng Wan
2025-12-28 14:57:04 -08:00
committed by GitHub
parent d7a3336ebe
commit 6f9d0a89a0
8 changed files with 84 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/bench_one_batch.py
View File

@@ -302,7 +302,7 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts):
)
req.fill_ids = req.origin_input_ids
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = -1
reqs.append(req)
return input_ids, reqs
@@ -318,7 +318,7 @@ def prepare_extend_inputs_for_correctness_test(
i, : bench_args.cut_len
]
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = -1
return reqs
@@ -345,7 +345,7 @@ def prepare_synthetic_inputs_for_latency_test(
)
req.fill_ids = req.origin_input_ids
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = -1
reqs.append(req)
return reqs

41
python/sglang/srt/managers/schedule_batch.py
View File

@@ -851,7 +851,7 @@ class Req:
input_len = len(self.fill_ids)
# NOTE: the matched length is at most 1 less than the input length to enable logprob computation
max_prefix_len = input_len - 1
if self.return_logprob:
if self.return_logprob and self.logprob_start_len >= 0:
max_prefix_len = min(max_prefix_len, self.logprob_start_len)
max_prefix_len = max(max_prefix_len, 0)
token_ids = self.fill_ids[:max_prefix_len]
@@ -1120,6 +1120,7 @@ class Req:
self.grammar = None
self.origin_input_ids = [0] # set it to one token to skip the long prefill
self.return_logprob = False
self.logprob_start_len = -1
self.to_finish = FINISH_ABORT(
error_msg, HTTPStatus.BAD_REQUEST, "BadRequestError"
)
@@ -1490,26 +1491,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# (= len(fill_ids) - len(prefix_indices), where fill_ids = origin_input_ids + output_ids
# and prefix_indices are the cached/shared prefix tokens)
#
if req.logprob_start_len >= pre_len:
# Optimization for prefill-only requests: When we only need logprobs at
# positions beyond the input sequence (to score next-token likelihood), skip all
# input logprob computation during prefill since no generation will occur.
if self.is_prefill_only and req.logprob_start_len == len(
req.origin_input_ids
):
# Skip ALL input logprobs: set extend_logprob_start_len = extend_input_len
req.extend_logprob_start_len = req.extend_input_len
else:
# Convert absolute logprob_start_len to relative extend_logprob_start_len
#
# Example: origin_input_ids=[1,2,3,4,5] (5 tokens, positions 0-4), logprob_start_len=3
# Regular logic: min(3-0, 5, 5-1) = min(3,5,4) = 3
# This means: "compute logprobs from position 3 onwards in extend batch"
req.extend_logprob_start_len = min(
req.logprob_start_len - pre_len,
req.extend_input_len,
req.seqlen - 1,
)
if req.logprob_start_len == -1:
req.extend_logprob_start_len = min(
len(req.fill_ids) - 1 - pre_len,
req.extend_input_len,
)
elif req.logprob_start_len >= pre_len:
req.extend_logprob_start_len = min(
req.logprob_start_len - pre_len,
req.extend_input_len,
)
else:
# logprob_start_len is before the current extend batch, so start from beginning
req.extend_logprob_start_len = 0
@@ -1532,9 +1523,13 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
len(req.prefix_indices),
len(req.fill_ids),
)
if req.logprob_start_len == -1:
logprob_start_len = len(req.origin_input_ids) - 1
else:
logprob_start_len = req.logprob_start_len
# Apply logprob_start_len
if global_start_idx < req.logprob_start_len:
global_start_idx = req.logprob_start_len
if global_start_idx < logprob_start_len:
global_start_idx = logprob_start_len
logprob_token_ids = req.origin_input_ids[
global_start_idx + 1 : global_end_idx + 1

23
python/sglang/srt/managers/scheduler.py
View File

@@ -1524,24 +1524,23 @@ class Scheduler(
self._add_request_to_queue(req)
return
# Copy more attributes
if recv_req.logprob_start_len == -1 or not recv_req.return_logprob:
# By default, only return the logprobs for output tokens
# For prefill-only requests with logprob_start_len == -1, set logprob_start_len beyond input sequence
# to skip input logprob computation entirely
if recv_req.logprob_start_len == -1:
if 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)
else:
# TODO: For text generation, evaluate setting logprob_start_len to len(req.origin_input_ids) as well
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
else:
req.logprob_start_len = recv_req.logprob_start_len
if not req.is_prefill_only and req.logprob_start_len >= len(
req.origin_input_ids
):
if req.logprob_start_len > len(req.origin_input_ids):
error_msg = f"{req.logprob_start_len=} is higher than the number of input tokens {len(req.origin_input_ids)=}. Please use a smaller logprob_start_len."
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = -1
req.set_finish_with_abort(error_msg)
self._add_request_to_queue(req)
return
@@ -1760,7 +1759,7 @@ class Scheduler(
return
# Copy more attributes
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = -1
self._add_request_to_queue(req)
def handle_batch_embedding_request(

22
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
View File

@@ -121,18 +121,18 @@ def prepare_mlp_sync_batch_raw(
num_tokens_for_logprob = num_tokens
else:
num_tokens = local_batch.extend_num_tokens
if local_batch.return_logprob:
num_tokens_for_logprob = sum(
# We should have at least 1 token for sample in every case.
max(extend_len - logprob_start_len, 1)
for logprob_start_len, extend_len in zip(
local_batch.extend_logprob_start_lens,
local_batch.extend_lens,
)
num_tokens_for_logprob = sum(
# We should have at least 1 token for sample in every case.
max(extend_len - logprob_start_len, 1)
for logprob_start_len, extend_len in zip(
local_batch.extend_logprob_start_lens,
local_batch.extend_lens,
)
else:
# When return_logprob = False, only need last token per request
num_tokens_for_logprob = local_batch.batch_size()
)
assert (
local_batch.return_logprob
or num_tokens_for_logprob == local_batch.batch_size()
)
skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get()
can_cuda_graph = (

4
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -591,10 +591,10 @@ class SchedulerOutputProcessorMixin:
For regular requests, all positions from logprob_start_len onwards have logprobs.
"""
is_multi_item_scoring = self._is_multi_item_scoring(req)
relevant_tokens = req.origin_input_ids[req.logprob_start_len :]
if is_multi_item_scoring:
# Multi-item scoring: count delimiter tokens from logprob_start_len onwards
relevant_tokens = req.origin_input_ids[req.logprob_start_len :]
return sum(
1
for token_id in relevant_tokens
@@ -602,7 +602,7 @@ class SchedulerOutputProcessorMixin:
)
else:
# Regular request: all tokens from logprob_start_len onwards
return len(req.origin_input_ids) - req.logprob_start_len
return len(relevant_tokens)
def _calculate_num_input_logprobs(
self, req: Req, extend_input_len: int, extend_logprob_start_len: int

2
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -565,7 +565,7 @@ class SchedulerPPMixin:
)
req.fill_ids = req.origin_input_ids
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = len(req.origin_input_ids) - 1
req.logprob_start_len = -1
# Prepare batch
batch = ScheduleBatch.init_new(