Refactor: Moving extend_logprob_start_len calculation out of prepare_for_extend (#16105)

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

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
 

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])

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)

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(

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)

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)