Fail fast on undersized swa pool (#37610)
This commit is contained in:
@@ -151,6 +151,27 @@ class MemoryPoolConfigurator:
|
|||||||
) -> MemoryPoolConfig:
|
) -> MemoryPoolConfig:
|
||||||
return config
|
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):
|
class DefaultPoolConfigurator(MemoryPoolConfigurator):
|
||||||
"""Configurator for standard models: MHA, MLA, DSA, FP4.
|
"""Configurator for standard models: MHA, MLA, DSA, FP4.
|
||||||
@@ -658,16 +679,9 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
|
|||||||
full_tokens = align_page_size(max_total_num_tokens)
|
full_tokens = align_page_size(max_total_num_tokens)
|
||||||
swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio))
|
swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio))
|
||||||
|
|
||||||
if (
|
self.validate_swa_pool_size(
|
||||||
self._sliding_window_size is not None
|
swa_tokens, self._sliding_window_size, self._page_size
|
||||||
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."
|
|
||||||
)
|
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
f"Use sliding window memory pool. "
|
f"Use sliding window memory pool. "
|
||||||
@@ -860,6 +874,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
|
|||||||
f"local={len(self.compression_ratios)}/{len(cfg.compress_ratios)}"
|
f"local={len(self.compression_ratios)}/{len(cfg.compress_ratios)}"
|
||||||
)
|
)
|
||||||
self.swa_page_size = cfg.window_size
|
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.swa_ratio = get_schedule().swa_full_tokens_ratio
|
||||||
self.is_speculative = get_spec().speculative_algorithm is not None
|
self.is_speculative = get_spec().speculative_algorithm is not None
|
||||||
self.online_c128_mtp_max_draft_tokens = max_speculative_num_draft_tokens() or 0
|
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:
|
def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes:
|
||||||
full_token = full_token // page_size * page_size
|
full_token = full_token // page_size * page_size
|
||||||
swa_tokens = int(full_token * self.swa_ratio) // 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(
|
return _DSV4PoolSizes(
|
||||||
full_max_total_num_tokens=full_token,
|
full_max_total_num_tokens=full_token,
|
||||||
swa_max_total_num_tokens=swa_tokens,
|
swa_max_total_num_tokens=swa_tokens,
|
||||||
|
|||||||
@@ -1007,5 +1007,66 @@ class TestDflashDraftKvBudget(CustomTestCase):
|
|||||||
self.assertLess(_tokens(10240), _tokens(None))
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user