fix(scheduler): keep SWA pages for the running batch's next decode

python/sglang/srt/managers/schedule_policy.py

@@ -657,6 +657,10 @@ def __init__(
         self.memory_budget = token_to_kv_pool_allocator.create_prefill_budget(
             tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
         )
+        if running_batch is not None and running_batch.reqs:
+            self.memory_budget.reserve_next_decode(
+                running_batch.new_tokens_required_next_decode()
+            )
 
         self.req_states = None
         self.can_run_list = []

python/sglang/srt/mem_cache/prefill_budget.py

@@ -150,6 +150,10 @@ def reserve(
         self.total_offset += immediate + max_new_tokens
         self.current_offset += immediate
 
+    def reserve_next_decode(self, num_tokens: int) -> None:
+        # total_offset already charges each running request's remaining decode.
+        pass
+
 
 class SWAPrefillBudget(PrefillBudget):
     """Separate FULL/SWA partitions, including per-request SWA rings."""
@@ -174,6 +178,12 @@ def remaining_swa(self):
         evictable = 0 if self.req_ring else self.tree_cache.swa_evictable_size()
         return self.allocator.swa_available_size() + evictable - self.swa_offset
 
+    def reserve_next_decode(self, num_tokens: int) -> None:
+        # Out-of-window SWA pages are freed only by decode, so a prefill that takes
+        # the next decode's pages makes it retract; a ring slot covers decode.
+        if not self.req_ring:
+            self.swa_offset += num_tokens
+
     def swa_tokens(
         self,
         extend_input_len,

test/registered/unit/managers/test_prefill_adder.py

@@ -831,6 +831,40 @@ def test_add_chunked_req_hybrid_swa_defers_when_swa_below_page(self):
         req.set_extend_range.assert_not_called()
         self.assertEqual(len(adder.can_run_list), 0)
 
+    def test_hybrid_swa_prefill_leaves_running_decode_pages(self):
+        """A prefill must not take the SWA pages the running batch's next decode
+        needs; otherwise that decode retracts requests."""
+        page, window = 64, 128
+        self.mock_tree_cache.sliding_window_size = window
+        self.mock_token_allocator.full_available_size.return_value = 100_000
+        self.mock_token_allocator.create_prefill_budget.side_effect = (
+            lambda tree_cache, **kwargs: SWAPrefillBudget(
+                self.mock_token_allocator, tree_cache, **kwargs
+            )
+        )
+        request = dict(
+            extend_input_len=256,
+            total_tokens=288,
+            max_new_tokens=32,
+            input_tokens=256,
+            swa_host_hit_length=0,
+            chunk_limit=None,
+        )
+        needed = estimate_swa_kv_tokens(
+            256, 32, sliding_window_size=window, page_size=page
+        )
+        self.mock_token_allocator.swa_available_size.return_value = needed + page
+        for decode_tokens, admitted in ((0, True), (page, False), (2 * page, False)):
+            with self.subTest(decode_tokens=decode_tokens):
+                running = self.create_running_batch(
+                    [self.create_mock_req("running", priority=0, max_new_tokens=8)]
+                )
+                running.new_tokens_required_next_decode.return_value = decode_tokens
+                adder = self.create_adder(running, page_size=page)
+                self.assertEqual(
+                    adder.memory_budget.check_prefill(**request)[0], admitted
+                )
+
     def test_swa_budget_for_req(self):
         # budget = max(alloc - window, 0) + min(extend + max_new, window) + page,
         # where alloc = min(extend, rem_chunk). The decode headroom is the SWA the