From 21225aba3d63000ba89b98fc9c85e4b6517c4b72 Mon Sep 17 00:00:00 2001 From: Peng Wu Date: Thu, 6 Aug 2026 07:09:04 -0700 Subject: [PATCH] [Scheduler] Fix to restrict the SWA chunk-cap escape hatch to true head-of-line livelock (#32700) --- python/sglang/srt/managers/schedule_policy.py | 33 ++++++++++++ test/manual/test_schedule_policy.py | 52 +++++++++++++++++++ .../unit/managers/test_prefill_adder.py | 2 + 3 files changed, 87 insertions(+) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 65526e595..4d79fc88e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 ) diff --git a/test/manual/test_schedule_policy.py b/test/manual/test_schedule_policy.py index c04481a46..0fff50e39 100644 --- a/test/manual/test_schedule_policy.py +++ b/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() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index baf3b74f6..33943b4a2 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/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: