[Scheduler] Fix to restrict the SWA chunk-cap escape hatch to true head-of-line livelock (#32700)

python/sglang/srt/managers/schedule_policy.py

@@ -704,6 +704,27 @@ class PrefillAdder:
             return 0
         return cap // self.page_size * self.page_size
 
+    def _swa_req_never_fits(
+        self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
+    ) -> bool:
+        """True when a request's SWA budget exceeds the *entire* SWA pool, so it
+        can never be admitted whole no matter how far the pool drains.
+
+        This is the head-of-line livelock the _swa_chunk_cap escape hatch exists
+        for; the hatch must fire only in this case. A request that merely
+        exceeds *current* rem_swa (transient pressure) would fit once running
+        decodes free their windows, so it must wait — admitting it into the
+        decode headroom collapses the SWA evictable cushion and forces running
+        requests to retract (observed as a severe retraction/re-prefill storm on
+        hybrid-SWA models at high concurrency)."""
+        capacity = self.token_to_kv_pool_allocator.size_swa
+        return (
+            self._swa_budget_for_req(
+                extend_input_len, max_new_tokens, swa_host_hit_length
+            )
+            >= capacity
+        )
+
     def _mamba_gap_budget_for_req(self, req: Req) -> int:
         """Shared-gap reservation (full-token-equivalents) for a request's new
         mamba state. Charged only on the SHARED Mamba pool (`_mamba_slot_cost > 0`)
@@ -1125,6 +1146,12 @@ class PrefillAdder:
                 swa_host_hit_length=req.swa_host_hit_length,
             )
             if swa_needed >= self.rem_swa_tokens:
+                if not self._swa_req_never_fits(
+                    real_input_tokens,
+                    self._swa_new_tokens(req),
+                    req.swa_host_hit_length,
+                ):
+                    return AddReqResult.NO_TOKEN
                 swa_cap = self._swa_chunk_cap(
                     self._swa_new_tokens(req), req.swa_host_hit_length
                 )
@@ -1155,6 +1182,12 @@ class PrefillAdder:
                     swa_host_hit_length=req.swa_host_hit_length,
                 )
                 if swa_needed >= self.rem_swa_tokens:
+                    if not self._swa_req_never_fits(
+                        real_input_tokens,
+                        self._swa_new_tokens(req),
+                        req.swa_host_hit_length,
+                    ):
+                        return AddReqResult.NO_TOKEN
                     swa_cap = self._swa_chunk_cap(
                         self._swa_new_tokens(req), req.swa_host_hit_length
                     )

test/manual/test_schedule_policy.py

@@ -1,10 +1,12 @@
 import unittest
 from array import array
+from types import SimpleNamespace
 
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.schedule_policy import (
     CacheAgnosticPolicy,
     CacheAwarePolicy,
+    PrefillAdder,
     SchedulePolicy,
 )
 from sglang.srt.mem_cache.radix_cache import RadixCache
@@ -327,5 +329,55 @@ class TestSchedulePolicy(CustomTestCase):
         self.assertEqual(waiting_queue[2].rid, "w3")
 
 
+def _swa_adder(size_swa, sliding_window, page_size=16, rem_chunk_tokens=512):
+    """A PrefillAdder carrying only the fields _swa_req_never_fits transitively
+    reads, built via __new__ so the check is tested as pure logic (no KV pools,
+    no GPU)."""
+    adder = PrefillAdder.__new__(PrefillAdder)
+    adder.page_size = page_size
+    adder.rem_chunk_tokens = rem_chunk_tokens
+    adder.tree_cache = SimpleNamespace(sliding_window_size=sliding_window)
+    adder.token_to_kv_pool_allocator = SimpleNamespace(size_swa=size_swa)
+    return adder
+
+
+class TestSwaChunkCapHatch(CustomTestCase):
+    """The _swa_chunk_cap shrink-and-admit hatch must fire ONLY for a request
+    whose SWA budget can never fit the drained pool (true head-of-line
+    livelock). _swa_req_never_fits() is the gate: True => hatch, False => wait
+    (NO_TOKEN). Admitting transient-pressure requests instead of waiting
+    collapses the SWA evictable cushion and causes a retraction storm."""
+
+    def test_transient_pressure_request_waits(self):
+        # Small request against an ample pool: budget << pool, so it would fit
+        # once running decodes drain -> must wait, not take the hatch.
+        adder = _swa_adder(size_swa=1024, sliding_window=128)
+        self.assertFalse(
+            adder._swa_req_never_fits(extend_input_len=256, max_new_tokens=64)
+        )
+
+    def test_request_larger_than_whole_pool_takes_hatch(self):
+        # A large host-hit load-back charge pushes the budget past the entire
+        # pool: it can never fit however far the pool drains -> hatch (True).
+        adder = _swa_adder(size_swa=1024, sliding_window=128)
+        self.assertTrue(
+            adder._swa_req_never_fits(
+                extend_input_len=256, max_new_tokens=64, swa_host_hit_length=4096
+            )
+        )
+
+    def test_decision_is_gated_by_pool_capacity(self):
+        # Same request; only the pool size changes. Proves the check compares
+        # the budget against size_swa (guards against a wrong-accessor bug):
+        # never-fits on a small pool, fits on a large one.
+        req = dict(extend_input_len=256, max_new_tokens=64, swa_host_hit_length=600)
+        self.assertTrue(
+            _swa_adder(size_swa=512, sliding_window=128)._swa_req_never_fits(**req)
+        )
+        self.assertFalse(
+            _swa_adder(size_swa=4096, sliding_window=128)._swa_req_never_fits(**req)
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/managers/test_prefill_adder.py

@@ -58,11 +58,13 @@ class TestPrefillAdder(CustomTestCase):
         full_available_size: int = 0,
         swa_available_size: int = 0,
         available_size: int = 0,
+        size_swa: int = 1_000_000,
     ) -> MagicMock:
         allocator = MagicMock()
         allocator.full_available_size.return_value = full_available_size
         allocator.swa_available_size.return_value = swa_available_size
         allocator.available_size.return_value = available_size
+        allocator.size_swa = size_swa
         return allocator
 
     def create_running_batch(self, reqs=None) -> MagicMock: