Avoid scattered assignment of extend_input_len and fill_len by merging them into Req.extend_range (#27610)

python/sglang/srt/disaggregation/decode.py

@@ -1420,7 +1420,6 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         # inserts committed KV into the radix tree. The last output token
         # hasn't had KV committed yet (output_ids is 1 ahead).
         req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids
-        req.fill_len = req.kv_committed_len
         # Set prefix_indices so downstream consumers (init_next_round_input,
         # prepare_for_extend) see the correct prefix length. In the agg path
         # this is done inside init_next_round_input, but decode-disagg needs
@@ -1428,7 +1427,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         req.prefix_indices = (
             prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64)
         )
-        req.set_extend_input_len(req.fill_len - total_prefix_len)
+        req.set_extend_range(total_prefix_len, req.kv_committed_len)
 
         # Return the transfer destination indices:
         if self.scheduler.enable_hisparse:
@@ -1909,8 +1908,7 @@ class SchedulerDisaggregationDecodeMixin:
                 # only sees committed KV (full array includes one uncommitted
                 # token because init_next_round_input rebuilt it as full).
                 if req.kv_committed_len is not None:
-                    req.fill_len = req.kv_committed_len
-                    req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
+                    req.set_extend_range(len(req.prefix_indices), req.kv_committed_len)
             else:
                 waiting_queue.append(req)
 

python/sglang/srt/dllm/mixin/req.py

@@ -57,7 +57,7 @@ class ReqDllmMixin:
     def _init_fill_ids_for_dllm(self: Req):
         self.dllm_block_offset = (
             0
-            if self.fill_len == 0
+            if not self.dllm_initialized
             else self.dllm_block_offset + self.dllm_config.block_size
         )
         self.full_untruncated_fill_ids = (
@@ -65,7 +65,7 @@ class ReqDllmMixin:
             + self.output_ids
             + array("q", [self.dllm_config.mask_id] * self.dllm_config.block_size)
         )
-        self.fill_len = len(self.full_untruncated_fill_ids)
+        self.dllm_initialized = True
 
     def _update_block_offset_for_dllm(self):
         prefix_len = len(self.prefix_indices)

python/sglang/srt/managers/schedule_batch.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.utils.common import (
+    Range,
     ceil_align,
     flatten_arrays_to_pinned_cpu,
     is_pin_memory_available,
@@ -726,7 +727,8 @@ class Req(ReqDllmMixin):
         # Kept in sync by _refresh_fill_ids; admission only updates fill_len,
         # never mutates this array's length.
         self.full_untruncated_fill_ids = array("q")
-        self.fill_len: int = 0
+        self.extend_range: Optional[Range] = None
+        self.dllm_initialized: bool = False
 
         self.session = session
         self.input_embeds = input_embeds
@@ -842,8 +844,6 @@ class Req(ReqDllmMixin):
         # Prefix info
         # The indices to kv cache for the shared prefix.
         self.prefix_indices: torch.Tensor = torch.empty((0,), dtype=torch.int64)
-        # Number of tokens to run prefill.
-        self.extend_input_len = 0
         # The relative logprob_start_len in an extend batch
         self.extend_logprob_start_len = 0
         # TODO(ispobock): rename to last_device_node
@@ -1096,6 +1096,18 @@ class Req(ReqDllmMixin):
         # Whether request reached finished condition
         return self.finished_reason is not None
 
+    @property
+    def fill_len(self) -> int:
+        return self.extend_range.end
+
+    @property
+    def extend_input_len(self) -> int:
+        return self.extend_range.length
+
+    def set_extend_range(self, start: int, end: int) -> None:
+        self.extend_range = Range(start, end)
+        self._recompute_extend_logprob_start_len()
+
     def get_fill_ids(self) -> array:
         return self.full_untruncated_fill_ids[: self.fill_len]
 
@@ -1447,7 +1459,8 @@ class Req(ReqDllmMixin):
         self.num_matched_prefix_tokens = 0
         self.swa_uuid_for_lock = None
         self.swa_prefix_lock_released = False
-        self.extend_input_len = 0
+        self.extend_range = None
+        self.dllm_initialized = False
         self.is_retracted = True
         self.retracted_stain = True
         self.input_token_logprobs = None
@@ -1470,7 +1483,6 @@ class Req(ReqDllmMixin):
         self.swa_evicted_seqlen = 0
         self.extend_batch_idx = 0
         self.decode_batch_idx = 0
-        self.fill_len = 0
 
         # When using input_embeds, we cannot easily mix the original input embeddings
         # with the newly generated output token IDs during re-prefill of retracted request.
@@ -1520,14 +1532,13 @@ class Req(ReqDllmMixin):
         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):
+    def _recompute_extend_logprob_start_len(self):
         # Setting extend_input_len and computing 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
-        self.extend_input_len = extend_input_len
         if self.logprob_start_len == -1:
             logprob_start_len = len(self.full_untruncated_fill_ids)
         else:
@@ -1997,7 +2008,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             if encoder_len == 0:
                 continue
             if len(req.prefix_indices) < encoder_len:
-                req.extend_input_len -= encoder_len
+                assert len(req.prefix_indices) == 0
+                req.extend_range = req.extend_range._replace(
+                    start=req.extend_range.start + encoder_len
+                )
                 req.extend_logprob_start_len = max(
                     0, req.extend_logprob_start_len - encoder_len
                 )
@@ -2382,8 +2396,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
         for req in running_batch.reqs:
             req._refresh_fill_ids()
-            req.fill_len = len(req.full_untruncated_fill_ids)
-            req.set_extend_input_len(1)
+            full_len = len(req.full_untruncated_fill_ids)
+            req.set_extend_range(full_len - 1, full_len)
 
         # Decode tokens of the running portion live in future_map.output_tokens_buf.
         self.input_ids = None

python/sglang/srt/managers/schedule_policy.py

@@ -659,8 +659,7 @@ class PrefillAdder:
             * self.page_size
         )
 
-        req.set_extend_input_len(trunc_len)
-        req.fill_len = prefix_len + trunc_len
+        req.set_extend_range(prefix_len, prefix_len + trunc_len)
 
         self.can_run_list.append(req)
 
@@ -684,8 +683,7 @@ class PrefillAdder:
         )
         truncated = cand_extend_input_len > _rem_tokens
         new_len = min(cand_extend_input_len, _rem_tokens)
-        req.set_extend_input_len(new_len)
-        req.fill_len = len(req.prefix_indices) + new_len
+        req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len)
         self.can_run_list.append(req)
 
         # Update budget: reserve max_new_tokens only if not truncated
@@ -728,8 +726,7 @@ class PrefillAdder:
         )
         truncated = cand_extend_input_len > _rem_tokens
         new_len = min(cand_extend_input_len, _rem_tokens)
-        req.set_extend_input_len(new_len)
-        req.fill_len = len(req.prefix_indices) + new_len
+        req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len)
         self.can_run_list.append(req)
         self._update_prefill_budget(
             0,
@@ -839,10 +836,9 @@ class PrefillAdder:
             or cand_extend_input_len <= self.rem_chunk_tokens  # it is the last chunk
         ):
             # Non-chunked prefill — the whole sequence is committed this iter.
-            req.set_extend_input_len(
-                len(req.full_untruncated_fill_ids) - len(req.prefix_indices)
+            req.set_extend_range(
+                len(req.prefix_indices), len(req.full_untruncated_fill_ids)
             )
-            req.fill_len = len(req.full_untruncated_fill_ids)
             self.can_run_list.append(req)
             self._update_prefill_budget(
                 0,
@@ -857,9 +853,10 @@ class PrefillAdder:
             # Chunked prefill
             trunc_len = self.rem_chunk_tokens
 
-            req.set_extend_input_len(trunc_len)
             assert len(req.prefix_indices) == 0
-            req.fill_len = len(req.prefix_indices) + trunc_len
+            req.set_extend_range(
+                len(req.prefix_indices), len(req.prefix_indices) + trunc_len
+            )
             self.can_run_list.append(req)
             self.new_chunked_req = req
             self._update_prefill_budget(0, trunc_len, 0, req.retracted_stain)
@@ -980,10 +977,9 @@ class PrefillAdder:
                 self._req_inc_lock_ref(req)
             elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
                 # Non-chunked prefill — the whole sequence is committed this iter.
-                req.set_extend_input_len(
-                    len(req.full_untruncated_fill_ids) - len(req.prefix_indices)
+                req.set_extend_range(
+                    len(req.prefix_indices), len(req.full_untruncated_fill_ids)
                 )
-                req.fill_len = len(req.full_untruncated_fill_ids)
                 self.can_run_list.append(req)
 
                 self._req_inc_lock_ref(req)
@@ -1022,8 +1018,9 @@ class PrefillAdder:
                     return AddReqResult.OTHER
 
                 # Chunked prefill
-                req.set_extend_input_len(trunc_len)
-                req.fill_len = len(req.prefix_indices) + trunc_len
+                req.set_extend_range(
+                    len(req.prefix_indices), len(req.prefix_indices) + trunc_len
+                )
 
                 self.can_run_list.append(req)
                 self.new_chunked_req = req

python/sglang/srt/managers/scheduler.py

@@ -3337,7 +3337,8 @@ class Scheduler(
             # we can use the correct values in output processing.
             if batch.return_logprob:
                 batch_result.extend_input_len_per_req = [
-                    req.extend_input_len for req in batch.reqs
+                    req.extend_input_len if req.extend_range is not None else 0
+                    for req in batch.reqs
                 ]
                 batch_result.extend_logprob_start_len_per_req = [
                     req.extend_logprob_start_len for req in batch.reqs

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -608,9 +608,10 @@ class SchedulerPPMixin:
                     sampling_params=sampling_params,
                 )
                 req.full_untruncated_fill_ids = req.origin_input_ids
-                req.fill_len = len(req.full_untruncated_fill_ids)
                 req.logprob_start_len = -1
-                req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
+                req.set_extend_range(
+                    len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+                )
 
                 # Prepare batch
                 batch = ScheduleBatch.init_new(

python/sglang/srt/utils/common.py

@@ -65,6 +65,7 @@ from typing import (
     Dict,
     Generic,
     List,
+    NamedTuple,
     Optional,
     Protocol,
     Sequence,
@@ -105,6 +106,15 @@ logger = logging.getLogger(__name__)
 torch_release = pkg_version.parse(torch.__version__).release
 
 
+class Range(NamedTuple):
+    start: int
+    end: int
+
+    @property
+    def length(self) -> int:
+        return self.end - self.start
+
+
 def flatten_arrays_to_pinned_cpu(parts: List[array[int]], pin: bool) -> torch.Tensor:
     """Flatten array.array('q') buffers into one int64 CPU tensor.
 

test/manual/test_forward_split_prefill.py

@@ -96,9 +96,10 @@ class TestForwardSplitPrefill(CustomTestCase):
                 sampling_params=sampling_params,
             )
             req.full_untruncated_fill_ids = req.origin_input_ids
-            req.fill_len = len(req.full_untruncated_fill_ids)
             req.logprob_start_len = -1
-            req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
+            req.set_extend_range(
+                len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+            )
             reqs.append(req)
 
         # Create dummy tree_cache for tests (no prefix caching, just allocation)

test/registered/unit/managers/test_hisparse_unit.py

@@ -15,8 +15,20 @@ from types import SimpleNamespace
 import torch
 
 from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
+from sglang.srt.utils.common import Range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
+
+class _FakeReq(SimpleNamespace):
+    @property
+    def fill_len(self) -> int:
+        return self.extend_range.end
+
+    @property
+    def extend_input_len(self) -> int:
+        return self.extend_range.length
+
+
 register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
 
@@ -42,7 +54,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
         origin_input_ids = list(range(64))
     if output_ids is None:
         output_ids = []
-    req = SimpleNamespace(
+    req = _FakeReq(
         rid=rid,
         origin_input_ids=origin_input_ids,
         output_ids=output_ids,
@@ -57,8 +69,8 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
         inflight_middle_chunks=0,
     )
     req.finished = lambda: req.finished_reason is not None
-    req.set_extend_input_len = lambda extend_input_len: setattr(
-        req, "extend_input_len", extend_input_len
+    req.set_extend_range = lambda start, end: setattr(
+        req, "extend_range", Range(start, end)
     )
     return req
 
@@ -220,7 +232,7 @@ class TestHiSparseUnit(unittest.TestCase):
         req.kv_allocated_len = fill_len
         req.kv_committed_len = fill_len
         req.full_untruncated_fill_ids = array("q", range(fill_len))
-        req.fill_len = fill_len
+        req.extend_range = Range(0, fill_len)
         return kv_loc
 
     # ==================================================================

test/registered/unit/managers/test_prefill_adder.py

@@ -76,7 +76,6 @@ class TestPrefillAdder(CustomTestCase):
         req = MagicMock(spec=Req)
         req.rid = str(rid)
         req.priority = priority
-        req.extend_input_len = 0
         req.prefix_indices = []
         req.full_untruncated_fill_ids = []
         req.extend_logprob_start_len = 0
@@ -385,11 +384,9 @@ class TestPrefillAdder(CustomTestCase):
 
         # Add a prefill that exactly consumes the chunk budget
         req1 = self.create_mock_req("req1", priority=0, max_new_tokens=64)
-        req1.extend_input_len = 56
         req1.host_hit_length = 0
         req1.prefix_indices = []
         req1.full_untruncated_fill_ids = list(range(56))
-        req1.fill_len = 56
         req1.last_node = MagicMock()
         req1.sampling_params.ignore_eos = False
 
@@ -420,11 +417,9 @@ class TestPrefillAdder(CustomTestCase):
 
         # Same prefill no longer exhausts the chunk budget
         req2 = self.create_mock_req("req2", priority=0, max_new_tokens=64)
-        req2.extend_input_len = 56
         req2.host_hit_length = 0
         req2.prefix_indices = []
         req2.full_untruncated_fill_ids = list(range(56))
-        req2.fill_len = 56
         req2.last_node = MagicMock()
         req2.sampling_params.ignore_eos = False
 
@@ -438,11 +433,9 @@ class TestPrefillAdder(CustomTestCase):
 
         # Fit last small prefill request
         req3 = self.create_mock_req("req3", priority=0, max_new_tokens=16)
-        req3.extend_input_len = 3
         req3.host_hit_length = 0
         req3.prefix_indices = []
         req3.full_untruncated_fill_ids = list(range(3))
-        req3.fill_len = 3
         req3.last_node = MagicMock()
         req3.sampling_params.ignore_eos = False
 
@@ -476,11 +469,9 @@ class TestPrefillAdder(CustomTestCase):
         adder.is_hybrid_swa = is_hybrid_swa
 
         req = self.create_mock_req("chunked", priority=0, max_new_tokens=128)
-        req.extend_input_len = extend_input_len
         req.prefix_indices = []
         req.full_untruncated_fill_ids = list(range(extend_input_len))
-        req.fill_len = extend_input_len
-        req.set_extend_input_len = MagicMock()
+        req.set_extend_range = MagicMock()
         return adder, req
 
     def test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend(self):
@@ -497,8 +488,9 @@ class TestPrefillAdder(CustomTestCase):
         result = adder.add_chunked_req(req)
 
         self.assertIs(result, req)  # truncated → chunked prefill continues
-        req.set_extend_input_len.assert_called_once()
-        new_len = req.set_extend_input_len.call_args.args[0]
+        req.set_extend_range.assert_called_once()
+        start, end = req.set_extend_range.call_args.args
+        new_len = end - start
         self.assertLessEqual(new_len + PAGE_SIZE, REM_SWA)
         self.assertEqual(new_len, REM_SWA - PAGE_SIZE)
 
@@ -510,13 +502,11 @@ class TestPrefillAdder(CustomTestCase):
         adder, req = self._build_hybrid_swa_chunked_req(
             page_size=PAGE_SIZE, rem_swa=PAGE_SIZE
         )
-        original_len = req.extend_input_len
 
         result = adder.add_chunked_req(req)
 
         self.assertIs(result, req)
-        req.set_extend_input_len.assert_not_called()
-        self.assertEqual(req.extend_input_len, original_len)
+        req.set_extend_range.assert_not_called()
         self.assertEqual(len(adder.can_run_list), 0)
 
     def test_swa_budget_for_req(self):
@@ -554,7 +544,7 @@ class TestPrefillAdder(CustomTestCase):
 
         result = adder.add_chunked_req(req)
         self.assertIsNone(result)
-        req.set_extend_input_len.assert_called_once_with(200)
+        req.set_extend_range.assert_called_once_with(0, 200)
         self.assertIn(req, adder.can_run_list)
 
 

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -15,6 +15,7 @@ maybe_stub_sgl_kernel()
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.managers.scheduler import Scheduler
 from sglang.srt.mem_cache.chunk_cache import ChunkCache
+from sglang.srt.utils.common import Range
 
 register_cpu_ci(est_time=6, suite="base-a-test-cpu")
 
@@ -32,10 +33,9 @@ def _make_req(
     req.origin_input_ids = array("q", fill_ids)
     req.output_ids = array("q")
     req.full_untruncated_fill_ids = array("q", fill_ids)
-    req.fill_len = fill_len
     req.prefix_indices = prefix_indices
     req.req_pool_idx = req_pool_idx
-    req.extend_input_len = extend_input_len
+    req.extend_range = Range(fill_len - extend_input_len, fill_len)
     req.inflight_middle_chunks = 0
     req.host_hit_length = 0
     req.cache_protected_len = 0

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -38,6 +38,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
 )
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
+from sglang.srt.utils.common import Range
 
 
 def _make_cache_with_pools(page_size=1):
@@ -68,7 +69,7 @@ class MockReq:
 
     def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
         self.full_untruncated_fill_ids = array("q", fill_ids)
-        self.fill_len = len(self.full_untruncated_fill_ids)
+        self.extend_range = Range(0, len(self.full_untruncated_fill_ids))
         self.origin_input_ids = array(
             "q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
         )
@@ -83,6 +84,10 @@ class MockReq:
         self.kv_allocated_len = len(fill_ids)
         self.kv_committed_freed = False
 
+    @property
+    def fill_len(self):
+        return self.extend_range.end
+
     def get_fill_ids(self):
         return self.full_untruncated_fill_ids[: self.fill_len]
 

test/registered/unit/mem_cache/test_unified_radix_cache_bench.py

@@ -639,7 +639,9 @@ def bench_cache_finished(
         req.origin_input_ids = array("q", seq)
         req.output_ids = array("q")
         req.full_untruncated_fill_ids = array("q", seq)
-        req.fill_len = len(req.full_untruncated_fill_ids)
+        req.set_extend_range(
+            len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+        )
         req.last_node = node
         req.cache_protected_len = matched_len
         req.kv_committed_len = len(seq)

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -898,7 +898,9 @@ class UnifiedRadixCacheSuite:
         req.swa_uuid_for_lock = None
         req.extra_key = None
         req.full_untruncated_fill_ids = array("q", input_ids + output_ids)
-        req.fill_len = len(req.full_untruncated_fill_ids)
+        req.set_extend_range(
+            len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+        )
         if self.cfg.has_mamba:
             req.mamba_last_track_seqlen = kv_len
 
@@ -922,7 +924,9 @@ class UnifiedRadixCacheSuite:
         req.origin_input_ids = array("q", prompt_ids)
         req.output_ids = array("q", output_ids)
         req.full_untruncated_fill_ids = array("q", prompt_ids + output_ids)
-        req.fill_len = len(req.full_untruncated_fill_ids)
+        req.set_extend_range(
+            len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+        )
         kv_len = req.fill_len
         kv_indices = self._alloc(allocator, kv_len)
         req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -977,7 +981,9 @@ class UnifiedRadixCacheSuite:
         req.swa_uuid_for_lock = None
         req.extra_key = None
         req.full_untruncated_fill_ids = array("q", tokens)
-        req.fill_len = len(req.full_untruncated_fill_ids)
+        req.set_extend_range(
+            len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+        )
 
         avail_before = allocator.available_size()
         tree.cache_finished_req(req, is_insert=False)
@@ -995,7 +1001,9 @@ class UnifiedRadixCacheSuite:
         req.origin_input_ids = array("q", tokens)
         req.output_ids = array("q")
         req.full_untruncated_fill_ids = array("q", tokens)
-        req.fill_len = len(req.full_untruncated_fill_ids)
+        req.set_extend_range(
+            len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+        )
         kv_len = len(tokens)
         kv_indices = self._alloc(allocator, kv_len)
         req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -1130,7 +1138,9 @@ class UnifiedRadixCacheSuite:
         req.swa_uuid_for_lock = None
         req.extra_key = None
         req.full_untruncated_fill_ids = array("q", input_ids)
-        req.fill_len = len(req.full_untruncated_fill_ids)
+        req.set_extend_range(
+            len(req.prefix_indices), len(req.full_untruncated_fill_ids)
+        )
         if self.cfg.has_mamba:
             req.mamba_last_track_seqlen = kv_len