diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index cb67a7755..701e0abc6 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 7e09b14e3..4298907da 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/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)