minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Cap SWA pool sizing with chunk cache (#28755)

Browse Source
This commit is contained in:
cctry
2026-06-21 01:06:59 -07:00
committed by GitHub
parent c9488241e9
commit 6d4ca9bc54
12 changed files with 398 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/model_executor/test_pool_configurator.py
+228 -3
View File
@@ -16,14 +16,22 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu")
@contextlib.contextmanager
def mock_cpu_env(kv_size=2, tp_size=1):
"""Mock GPU-dependent functions for CPU-only testing."""
def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4):
"""Mock GPU-dependent functions for CPU-only testing.
swa_eviction_interval pins SGLANG_SWA_EVICTION_INTERVAL (decode batches between
SWA evictions) to a small value so the chunk-cap formula stays hand-computable;
only SWAChunkCapPoolConfigurator reads it.
"""
from sglang.srt.environ import envs
with (
patch("torch._utils._element_size", return_value=kv_size),
patch(
"sglang.srt.model_executor.pool_configurator.get_attention_tp_size",
return_value=tp_size,
),
envs.SGLANG_SWA_EVICTION_INTERVAL.override(swa_eviction_interval),
):
yield
@@ -44,6 +52,18 @@ def _make_model_runner(
swa_full_tokens_ratio=0.5,
page_size=1,
mambaish_config=None,
disable_radix_cache=False,
chunked_prefill_size=None,
disable_overlap_schedule=False,
sliding_window_size=None,
speculative_num_draft_tokens=None,
max_speculative_num_draft_tokens=None,
speculative_algorithm=None,
speculative_num_steps=None,
speculative_eagle_topk=None,
disaggregation_mode="null",
max_running_requests=None,
disaggregation_decode_extra_slots=0,
):
"""Create a mock ModelRunner with the fields configurators need."""
mr = MagicMock()
@@ -53,8 +73,11 @@ def _make_model_runner(
mr.num_effective_layers = num_layers
mr.start_layer = 0
mr.end_layer = num_layers
mr.dp_size = 1
mr.page_size = page_size
mr.mambaish_config = mambaish_config
mr.is_hybrid_swa = is_hybrid_swa
mr.sliding_window_size = sliding_window_size
mc = SimpleNamespace()
mc.head_dim = head_dim
@@ -80,6 +103,19 @@ def _make_model_runner(
sa = SimpleNamespace()
sa.swa_full_tokens_ratio = swa_full_tokens_ratio
sa.page_size = page_size
sa.disable_radix_cache = disable_radix_cache
sa.chunked_prefill_size = chunked_prefill_size
sa.disable_overlap_schedule = disable_overlap_schedule
sa.speculative_num_draft_tokens = speculative_num_draft_tokens
sa.max_speculative_num_draft_tokens = (
max_speculative_num_draft_tokens or speculative_num_draft_tokens
)
sa.speculative_algorithm = speculative_algorithm
sa.speculative_num_steps = speculative_num_steps
sa.speculative_eagle_topk = speculative_eagle_topk
sa.disaggregation_mode = disaggregation_mode
sa.max_running_requests = max_running_requests
sa.disaggregation_decode_extra_slots = disaggregation_decode_extra_slots
mr.server_args = sa
spec = MagicMock()
@@ -257,11 +293,172 @@ class TestHybridSWAConfigurator(unittest.TestCase):
int(config.full_max_total_num_tokens * 0.5),
)
def test_chunk_cache_cap_accounts_for_spec_topk_page_rounding(self):
available = 1_000_000
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
swa_num_kv_heads=4,
swa_full_tokens_ratio=0.5,
disable_radix_cache=True,
chunked_prefill_size=4,
sliding_window_size=8,
page_size=4,
max_running_requests=2,
speculative_algorithm="EAGLE",
speculative_num_steps=3,
speculative_eagle_topk=2,
speculative_num_draft_tokens=5,
disable_overlap_schedule=True, # spec-v1: no double allocation
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size=4)
# spec-v1 (overlap off): decode_alloc = max(ceil_align(3+4,4)*2,
# ceil_align(5,4)) = 16. trailing = 8 + 20 + page(4) = 32; per req =
# 32 + 16 = 48. Global prefill = 1*chunk(4) + page(4) = 8.
# cap = 48 * 2 + 8 = 104 -> ceil_align(104, 4) = 104.
self.assertEqual(config.swa_max_total_num_tokens, 104)
self.assertLessEqual(_actual_memory_used(mr, config), available)
def test_chunk_cache_cap_doubles_decode_alloc_for_spec_v2_overlap(self):
# Overlap on -> spec-v2: decode_alloc = 2 * get_alloc_len_per_decode =
# 2 * max(steps*topk, max_draft) = 2 * max(6, 5) = 12 (page=1, since the
# v2 allocator does not support page>1 & topk>1). trailing = 8 + 20 +
# page(1) = 29; per req = 29 + 12 = 41. Global prefill =
# 2*chunk(4) + page(1) = 9; cap = 41 * 2 + 9 = 91.
available = 1_000_000
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
swa_num_kv_heads=4,
swa_full_tokens_ratio=0.5,
disable_radix_cache=True,
chunked_prefill_size=4,
sliding_window_size=8,
page_size=1,
max_running_requests=2,
speculative_algorithm="EAGLE",
speculative_num_steps=3,
speculative_eagle_topk=2,
speculative_num_draft_tokens=5,
disable_overlap_schedule=False, # spec-v2: 2 * get_alloc_len_per_decode
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size=1)
self.assertEqual(config.swa_max_total_num_tokens, 91)
self.assertLessEqual(_actual_memory_used(mr, config), available)
def test_chunk_cache_cap_drops_prefill_for_disagg_decode(self):
available = 1_000_000
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
swa_num_kv_heads=4,
swa_full_tokens_ratio=0.5,
disable_radix_cache=True,
chunked_prefill_size=1000,
sliding_window_size=4,
page_size=1,
max_running_requests=10,
disaggregation_mode="decode",
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size=1)
# disagg decode drops the prefill term: per req = 4 + 1 + 4 + 1 = 10 (as above).
self.assertEqual(config.swa_max_total_num_tokens, 100)
self.assertLessEqual(_actual_memory_used(mr, config), available)
def test_chunk_cache_cap_prefill_holds_window_plus_chunk(self):
# Non-decode (prefill) engine: each request keeps its decode footprint, while
# in-flight chunked-prefill tokens are a global batch budget -- two chunks
# under overlap.
available = 1_000_000
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
swa_num_kv_heads=4,
swa_full_tokens_ratio=0.5,
disable_radix_cache=True,
chunked_prefill_size=16,
sliding_window_size=8,
page_size=4,
max_running_requests=2,
disaggregation_mode="prefill",
disable_overlap_schedule=False, # overlap -> 2 chunks in flight
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size=4)
# per req = trailing(window(8) + eviction(4) + page(4)) + decode_alloc(4)
# = 20. Global prefill = 2*chunk(16) + page(4) = 36.
# cap = 20 * max_running_requests(2) + 36 = 76.
self.assertEqual(config.swa_max_total_num_tokens, 76)
self.assertLessEqual(_actual_memory_used(mr, config), available)
def test_chunk_cache_cap_disagg_decode_pre_alloc(self):
# decode adds disaggregation_decode_extra_slots in-transfer slots to the
# request count (num_reserved_decode_tokens is a full-pool concern, not SWA).
available = 2_000_000
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
swa_num_kv_heads=4,
swa_full_tokens_ratio=0.5,
disable_radix_cache=True,
chunked_prefill_size=1000,
sliding_window_size=4,
page_size=1,
max_running_requests=10,
disaggregation_mode="decode",
disaggregation_decode_extra_slots=2,
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size=1)
# active per req = 4 + 1 + 4 + 1 = 10 for the 10 running requests; the 2
# in-transfer extra slots hold only window + page = 4 + 1 = 5 each.
# cap = 10 * 10 + 5 * 2 = 110.
self.assertEqual(config.swa_max_total_num_tokens, 110)
self.assertLessEqual(_actual_memory_used(mr, config), available)
class TestAllSWAConfigurator(unittest.TestCase):
"""All-SWA (full_layers=0): special case."""
def _run(self, available_bytes, ratio=0.5, page_size=1):
def _run(self, available_bytes, ratio=0.5, page_size=1, **kwargs):
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[],
@@ -269,6 +466,7 @@ class TestAllSWAConfigurator(unittest.TestCase):
swa_num_kv_heads=4,
swa_full_tokens_ratio=ratio,
page_size=page_size,
**kwargs,
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
@@ -362,6 +560,33 @@ class TestFactory(unittest.TestCase):
cfg = create_memory_pool_configurator(mr)
self.assertIsInstance(cfg, HybridSWAPoolConfigurator)
def test_chunk_cap_configurator_selection(self):
# SWAChunkCapPoolConfigurator is selected only when max_running_requests is set.
def _cfg(max_running_requests):
mr = _make_model_runner(
is_hybrid_swa=True,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
swa_num_kv_heads=4,
disable_radix_cache=True,
chunked_prefill_size=4,
sliding_window_size=8,
max_running_requests=max_running_requests,
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
return create_memory_pool_configurator(mr)
from sglang.srt.model_executor.pool_configurator import (
SWAChunkCapPoolConfigurator,
)
self.assertIsInstance(_cfg(2), SWAChunkCapPoolConfigurator)
self.assertNotIsInstance(_cfg(None), SWAChunkCapPoolConfigurator)
if __name__ == "__main__":
unittest.main()