Fix the _chunked_req_scheduled_last_iter flag with a content-based stash gate (#26938)

python/sglang/srt/managers/scheduler.py

@@ -936,15 +936,6 @@ class Scheduler(
         elif self.chunked_prefill_size is not None and self.chunked_prefill_size <= 0:
             self.chunked_prefill_size = None
         self.chunked_req = None
-        # Tracks whether the current self.chunked_req was actually scheduled
-        # into last iteration's batch (i.e., in can_run_list -> got a fresh
-        # req_pool_idx from prepare_for_extend). Used to gate the
-        # stash_chunked_request call at the top of get_next_batch_to_run:
-        # if add_chunked_req early-returned under hybrid-SWA pressure,
-        # the req_pool_idx was already freed and the full_untruncated_fill_ids
-        # was rebuilt by init_next_round_input, so running stash would
-        # double-free and corrupt prefix_indices.
-        self._chunked_req_scheduled_last_iter = False
         self.is_mixed_chunk = (
             self.chunked_prefill_size is not None
             and self.server_args.enable_mixed_chunk
@@ -2443,7 +2434,11 @@ class Scheduler(
             # only finished requests to running_batch.
             chunked_req_to_exclude.add(self.chunked_req)
 
-            if self._chunked_req_scheduled_last_iter:
+            # Stash (cache) the previous chunk only when it produced new KV
+            # beyond what is already cached. A parked chunk (add_chunked_req
+            # hybrid-SWA early-return) leaves fill_len == len(prefix_indices),
+            # so there is nothing new to cache and stashing would be a no-op.
+            if self.chunked_req.fill_len > len(self.chunked_req.prefix_indices):
                 self.stash_chunked_request(self.chunked_req)
 
         # HiSparse has its own prefill-to-decode transition; skip last_batch merge.
@@ -2645,11 +2640,6 @@ class Scheduler(
         if self.chunked_req is not None:
             self.chunked_req.init_next_round_input()
             self.chunked_req = adder.add_chunked_req(self.chunked_req)
-            self._chunked_req_scheduled_last_iter = (
-                self.chunked_req in adder.can_run_list
-            )
-        else:
-            self._chunked_req_scheduled_last_iter = False
 
         if self.enable_lora:
             running_loras = {req.lora_id for req in self.running_batch.reqs}
@@ -2749,9 +2739,6 @@ class Scheduler(
             # Update chunked prefill
             assert self.chunked_req is None
             self.chunked_req = adder.new_chunked_req
-            # new_chunked_req is added to can_run_list by add_one_req,
-            # so it will be scheduled this iter -> stash is needed next iter.
-            self._chunked_req_scheduled_last_iter = True
 
         if self.chunked_req is not None:
             self.chunked_req.inflight_middle_chunks += 1

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -25,13 +25,14 @@ def _make_req(
     fill_ids: list,
     prefix_indices: torch.Tensor,
     extend_input_len: int,
+    fill_len: int,
 ) -> Req:
     req = Req.__new__(Req)
     req.rid = "test-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 = len(req.full_untruncated_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
@@ -104,15 +105,17 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
 class TestStashGatePreservesPrefixIndices(CustomTestCase):
     """Consumer side: real ChunkCache.cache_unfinished_req mutates
     req.prefix_indices iff stash actually runs, so prefix_indices content
-    is the bug-detection signal."""
+    is the bug-detection signal. The stash gate is content-based:
+    `fill_len > len(prefix_indices)` means there is freshly computed KV to
+    cache; otherwise the chunk was parked and stashing must be skipped."""
 
     POOL_IDX = 4
     INITIAL_PREFIX_LEN = 8  # what was really cached last iter
-    POST_RESET_FILL_LEN = 32  # length after init_next_round_input
+    POST_RESET_FILL_LEN = 32  # length after init_next_round_input rebuilds
     NUM_SLOTS = 8
     MAX_CONTEXT = 64
 
-    def _build(self, flag: bool):
+    def _build(self, *, fill_len: int):
         pool = _make_req_to_token_pool(self.NUM_SLOTS, self.MAX_CONTEXT)
         cache = _make_chunk_cache(pool)
         initial_prefix = pool.req_to_token[self.POOL_IDX, : self.INITIAL_PREFIX_LEN].to(
@@ -122,16 +125,16 @@ class TestStashGatePreservesPrefixIndices(CustomTestCase):
             req_pool_idx=self.POOL_IDX,
             fill_ids=list(range(self.POST_RESET_FILL_LEN)),
             prefix_indices=initial_prefix,
-            extend_input_len=0,
+            extend_input_len=fill_len - self.INITIAL_PREFIX_LEN,
+            fill_len=fill_len,
         )
         s = _scheduler_for_get_next_batch(tree_cache=cache, chunked_req=req)
-        s._chunked_req_scheduled_last_iter = flag
         return s, req, initial_prefix, pool
 
-    def test_deferred_chunked_req_keeps_real_prefix_indices(self):
-        # The bug case: a spurious stash on a deferred chunked_req
-        # would extend prefix_indices to len(fill_ids).
-        s, req, initial_prefix, _ = self._build(flag=False)
+    def test_parked_chunked_req_keeps_real_prefix_indices(self):
+        # A parked chunk has fill_len == len(prefix_indices): no new KV was
+        # computed, so the gate must skip stash and leave prefix_indices intact.
+        s, req, initial_prefix, _ = self._build(fill_len=self.INITIAL_PREFIX_LEN)
 
         Scheduler.get_next_batch_to_run(s)
 
@@ -139,9 +142,9 @@ class TestStashGatePreservesPrefixIndices(CustomTestCase):
         self.assertTrue(torch.equal(req.prefix_indices, initial_prefix))
 
     def test_scheduled_chunked_req_advances_prefix_indices_via_real_stash(self):
-        # Symmetric guard against over-gating: when the chunked_req was
-        # actually scheduled, stash must run and advance prefix_indices.
-        s, req, _, pool = self._build(flag=True)
+        # Symmetric guard against over-gating: when fill_len has advanced past
+        # the cached prefix, stash must run and advance prefix_indices.
+        s, req, _, pool = self._build(fill_len=self.POST_RESET_FILL_LEN)
 
         Scheduler.get_next_batch_to_run(s)
 
@@ -151,13 +154,12 @@ class TestStashGatePreservesPrefixIndices(CustomTestCase):
         self.assertEqual(req.prefix_indices.shape[0], self.POST_RESET_FILL_LEN)
         self.assertTrue(torch.equal(req.prefix_indices, expected))
 
-    def test_no_chunked_req_never_mutates_state_even_with_stale_flag(self):
-        # Retract path clears chunked_req without resetting the flag;
-        # the outer `if chunked_req is not None` guard must hold.
+    def test_no_chunked_req_never_mutates_state(self):
+        # The outer `if chunked_req is not None` guard must hold on the retract
+        # path that clears chunked_req.
         pool = _make_req_to_token_pool(self.NUM_SLOTS, self.MAX_CONTEXT)
         cache = _make_chunk_cache(pool)
         s = _scheduler_for_get_next_batch(tree_cache=cache, chunked_req=None)
-        s._chunked_req_scheduled_last_iter = True
 
         Scheduler.get_next_batch_to_run(s)
         self.assertIsNone(s.chunked_req)