From ae3205ba284fff759f8001098559ceaae140af93 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sat, 5 Sep 2026 13:40:37 +0800 Subject: [PATCH] Fail fast on undersized swa pool (#37610) --- .../srt/model_executor/pool_configurator.py | 36 ++++++++--- .../model_executor/test_pool_configurator.py | 61 +++++++++++++++++++ 2 files changed, 87 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index dac5dd010..3a5e7d3da 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -151,6 +151,27 @@ class MemoryPoolConfigurator: ) -> MemoryPoolConfig: return config + @staticmethod + def validate_swa_pool_size( + swa_tokens: int, sliding_window_size: Optional[int], page_size: int + ) -> None: + """Reject an SWA pool too small to ever admit a request. + + Prefill charges min(extend + decode, window) + page_size of SWA headroom + per request, so a pool at or below that floor rejects every request no + matter how far it drains: the scheduler spins in the waiting queue and + the server hangs at warmup instead of failing here. + """ + if sliding_window_size is None: + return + if sliding_window_size + page_size >= swa_tokens: + raise ValueError( + f"SWA pool ({swa_tokens} tokens) cannot hold even one request: " + f"the prefill admission floor is sliding_window_size " + f"({sliding_window_size}) + page_size ({page_size}). " + f"Increase --swa-full-tokens-ratio or the total KV budget." + ) + class DefaultPoolConfigurator(MemoryPoolConfigurator): """Configurator for standard models: MHA, MLA, DSA, FP4. @@ -658,16 +679,9 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): full_tokens = align_page_size(max_total_num_tokens) swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio)) - if ( - self._sliding_window_size is not None - and self._sliding_window_size + self._page_size >= swa_tokens - ): - raise ValueError( - f"SWA pool ({swa_tokens} tokens) cannot hold even one request: " - f"the prefill admission floor is sliding_window_size " - f"({self._sliding_window_size}) + page_size ({self._page_size}). " - f"Increase --swa-full-tokens-ratio or the total KV budget." - ) + self.validate_swa_pool_size( + swa_tokens, self._sliding_window_size, self._page_size + ) logger.info( f"Use sliding window memory pool. " @@ -860,6 +874,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): f"local={len(self.compression_ratios)}/{len(cfg.compress_ratios)}" ) self.swa_page_size = cfg.window_size + self.sliding_window_size = kvc.sliding_window_size self.swa_ratio = get_schedule().swa_full_tokens_ratio self.is_speculative = get_spec().speculative_algorithm is not None self.online_c128_mtp_max_draft_tokens = max_speculative_num_draft_tokens() or 0 @@ -995,6 +1010,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes: full_token = full_token // page_size * page_size swa_tokens = int(full_token * self.swa_ratio) // page_size * page_size + self.validate_swa_pool_size(swa_tokens, self.sliding_window_size, page_size) return _DSV4PoolSizes( full_max_total_num_tokens=full_token, swa_max_total_num_tokens=swa_tokens, diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 95d73e6af..b76f5420d 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1007,5 +1007,66 @@ class TestDflashDraftKvBudget(CustomTestCase): self.assertLess(_tokens(10240), _tokens(None)) +class TestSWAPoolFloor(CustomTestCase): + """An SWA pool below the prefill admission floor must fail here, not livelock + the scheduler at warmup.""" + + def _hybrid_swa_from_max_tokens(self, max_tokens, ratio, page_size, window): + mr = _make_model_runner( + self, + is_hybrid_swa=True, + full_attention_layer_ids=list(range(16)), + swa_attention_layer_ids=list(range(16, 32)), + swa_num_kv_heads=4, + page_size=page_size, + swa_full_tokens_ratio=ratio, + sliding_window_size=window, + ) + with mock_cpu_env(): + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + cfg = create_memory_pool_configurator(mr) + return cfg.calculate_pool_sizes_from_max_tokens(max_tokens, page_size) + + def test_hybrid_swa_rejects_single_page_pool(self): + with self.assertRaisesRegex(ValueError, "cannot hold even one request"): + self._hybrid_swa_from_max_tokens( + max_tokens=4096, ratio=0.1, page_size=256, window=128 + ) + + def test_hybrid_swa_accepts_pool_above_floor(self): + config = self._hybrid_swa_from_max_tokens( + max_tokens=32768, ratio=0.1, page_size=256, window=128 + ) + self.assertEqual(config.swa_max_total_num_tokens, 3072) + + def _dsv4_sizes(self, max_tokens, page_size): + """Exercise the DSV4 size arithmetic without a full V4 model fixture: + _compute_dsv4_sizes reads only these five attributes.""" + from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator + + cfg = object.__new__(DSV4PoolConfigurator) + cfg.swa_ratio = 0.1 + cfg.sliding_window_size = 128 + cfg.swa_page_size = 128 + cfg.c4_ring_size = 8 + cfg.c4_shrink_factor = 1 + return cfg._compute_dsv4_sizes(max_tokens, page_size) + + def test_dsv4_rejects_single_page_pool(self): + # DeepSeek-V4-Flash defaults: page_size=256, swa_full_tokens_ratio=0.1. + # int(4096 * 0.1) page-aligns down to 256 -- exactly one page, below the + # 128 + 256 floor. + with self.assertRaisesRegex(ValueError, "cannot hold even one request"): + self._dsv4_sizes(max_tokens=4096, page_size=256) + + def test_dsv4_accepts_pool_above_floor(self): + sizes = self._dsv4_sizes(max_tokens=32768, page_size=256) + self.assertEqual(sizes.full_max_total_num_tokens, 32768) + self.assertEqual(sizes.swa_max_total_num_tokens, 3072) + + if __name__ == "__main__": unittest.main()