diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 692237326..b3163fbdf 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -301,8 +301,8 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts): sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) req.logprob_start_len = -1 + req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) reqs.append(req) return input_ids, reqs @@ -312,13 +312,13 @@ def prepare_extend_inputs_for_correctness_test( bench_args, input_ids, reqs, model_runner ): for i in range(len(reqs)): - req = reqs[i] + req: Req = reqs[i] req.fill_ids += input_ids[i][bench_args.cut_len :] req.prefix_indices = model_runner.req_to_token_pool.req_to_token[ i, : bench_args.cut_len ] - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) req.logprob_start_len = -1 + req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) return reqs @@ -344,8 +344,8 @@ def prepare_synthetic_inputs_for_latency_test( sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) req.logprob_start_len = -1 + req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) reqs.append(req) return reqs diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f4ef07e41..7726f6b64 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -690,7 +690,7 @@ class DecodePreallocQueue: # populate metadata req.fill_ids = req.origin_input_ids + req.output_ids - req.extend_input_len = len(req.fill_ids) + req.set_extend_input_len(len(req.fill_ids)) return kv_loc diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e49ad99f6..9263b9c3b 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -895,7 +895,7 @@ class Req: ) ) - self.extend_input_len = len(self.fill_ids) - len(self.prefix_indices) + self.set_extend_input_len(len(self.fill_ids) - len(self.prefix_indices)) # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313 def init_incremental_detokenize(self): @@ -1113,6 +1113,17 @@ class Req: logger.info(f"{prefix}: {self.time_stats.convert_to_duration()}") self.has_log_time_stats = True + def set_extend_input_len(self, extend_input_len: int): + self.extend_input_len = extend_input_len + if self.logprob_start_len == -1: + logprob_start_len = len(self.fill_ids) - 1 + else: + logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices)) + self.extend_logprob_start_len = min( + logprob_start_len - len(self.prefix_indices), + self.extend_input_len, + ) + def set_finish_with_abort(self, error_msg: str): if get_tensor_model_parallel_rank() == 0: logger.error(f"{error_msg}, {self.rid=}") @@ -1482,29 +1493,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_track_seqlens_cpu, ) - # Compute the relative logprob_start_len in an extend batch - # - # Key variables: - # - logprob_start_len: Absolute position in full sequence where logprob computation begins - # - extend_logprob_start_len: Relative position within current extend batch where logprob computation begins - # - extend_input_len: Number of tokens that need to be processed in this extend batch - # (= 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 == -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 - if self.return_logprob: # Find input logprob token ids. # First, find a global index within origin_input_ids and slide it by 1 @@ -1679,7 +1667,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): for req in running_batch.reqs: req.fill_ids = req.origin_input_ids + req.output_ids - req.extend_input_len = 1 + req.set_extend_input_len(1) input_ids = torch.cat([self.input_ids, running_batch.input_ids]) out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc]) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 32c0215c2..f88c002be 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -454,7 +454,7 @@ class PrefillAdder: if _rem_tokens <= 0: _rem_tokens = self.rem_chunk_tokens truncated = req.extend_input_len > _rem_tokens - req.extend_input_len = min(req.extend_input_len, _rem_tokens) + req.set_extend_input_len(min(req.extend_input_len, _rem_tokens)) req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len] self.can_run_list.append(req) self._update_prefill_budget( @@ -559,7 +559,7 @@ class PrefillAdder: # Chunked prefill trunc_len = self.rem_chunk_tokens - req.extend_input_len = trunc_len + req.set_extend_input_len(trunc_len) req.fill_ids = req.fill_ids[:trunc_len] self.can_run_list.append(req) self.new_chunked_req = req @@ -608,7 +608,7 @@ class PrefillAdder: req.last_host_node, req.host_hit_length ) req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) prefix_len = len(req.prefix_indices) req.cache_protected_len = prefix_len @@ -651,7 +651,7 @@ class PrefillAdder: ) # Chunked prefill - req.extend_input_len = trunc_len + req.set_extend_input_len(trunc_len) req.fill_ids = req.fill_ids[: len(req.prefix_indices) + trunc_len] self.can_run_list.append(req) diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 4e8af669f..5bc6f18be 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -564,8 +564,8 @@ class SchedulerPPMixin: sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) req.logprob_start_len = -1 + req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) # Prepare batch batch = ScheduleBatch.init_new( diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index 660289091..cf3ecb1c1 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/test/manual/test_forward_split_prefill.py @@ -92,8 +92,8 @@ class TestForwardSplitPrefill(CustomTestCase): sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) req.logprob_start_len = -1 + req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) reqs.append(req) # Create dummy tree_cache for tests (no prefix caching, just allocation) diff --git a/test/srt/test_prefill_adder.py b/test/srt/test_prefill_adder.py index a7308a401..60a8b904d 100644 --- a/test/srt/test_prefill_adder.py +++ b/test/srt/test_prefill_adder.py @@ -68,6 +68,7 @@ class TestPrefillAdder(CustomTestCase): req.rid = str(rid) req.priority = priority req.extend_input_len = 0 + req.extend_logprob_start_len = 0 req.output_ids = [0] * output_len req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens) req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time)