From 0756a1d2b070bb98f6c788c4c4c7346c68da164a Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Fri, 7 Aug 2026 17:49:07 +0800 Subject: [PATCH] Move SWA chunk-cap hatch tests into the registered suite (#33975) --- test/manual/test_schedule_policy.py | 52 ------------------- .../unit/managers/test_prefill_adder.py | 50 ++++++++++++++++++ 2 files changed, 50 insertions(+), 52 deletions(-) diff --git a/test/manual/test_schedule_policy.py b/test/manual/test_schedule_policy.py index 0fff50e39..c04481a46 100644 --- a/test/manual/test_schedule_policy.py +++ b/test/manual/test_schedule_policy.py @@ -1,12 +1,10 @@ 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 @@ -329,55 +327,5 @@ 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 8c63e45e9..8710ce8f5 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -840,6 +840,56 @@ class TestPrefillAdder(CustomTestCase): ): self.assertIsNone(adder._check_prefill_tile_budget(129)) + # ---- _swa_req_never_fits: the gate on the _swa_chunk_cap escape hatch ---- + # The hatch shrinks a chunk and admits it instead of waiting; it must fire + # ONLY for a request whose SWA budget can never fit the drained pool (true + # head-of-line livelock). Admitting transient-pressure requests instead + # collapses the SWA evictable cushion and causes a retraction storm. + + def create_swa_adder( + self, *, size_swa: int, sliding_window: int, page_size: int = 16 + ) -> PrefillAdder: + self.mock_tree_cache.sliding_window_size = sliding_window + self.mock_token_allocator = self.create_token_allocator(size_swa=size_swa) + return self.create_adder( + self.create_running_batch(), + page_size=page_size, + rem_chunk_tokens=512, + ) + + def test_swa_never_fits_false_under_transient_pressure(self): + # Small request against an ample pool: budget << pool, so it would fit + # once running decodes drain -> must wait, not take the hatch. + adder = self.create_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_swa_never_fits_true_when_budget_exceeds_whole_pool(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. + adder = self.create_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_swa_never_fits_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). + req = dict(extend_input_len=256, max_new_tokens=64, swa_host_hit_length=600) + self.assertTrue( + self.create_swa_adder(size_swa=512, sliding_window=128)._swa_req_never_fits( + **req + ) + ) + self.assertFalse( + self.create_swa_adder( + size_swa=4096, sliding_window=128 + )._swa_req_never_fits(**req) + ) + if __name__ == "__main__": unittest.main()