From 6d4ca9bc54d9181f1558119a9d0ae565ae2d66d3 Mon Sep 17 00:00:00 2001 From: cctry Date: Sun, 21 Jun 2026 01:06:59 -0700 Subject: [PATCH] Cap SWA pool sizing with chunk cache (#28755) --- .../ascend-npus/ascend_npu_best_practice.mdx | 6 +- .../best_practice/minimax_m2_5.mdx | 4 +- .../srt/arg_groups/pd_disaggregation_hook.py | 14 ++ python/sglang/srt/disaggregation/decode.py | 5 +- python/sglang/srt/environ.py | 7 +- python/sglang/srt/managers/schedule_batch.py | 17 +- python/sglang/srt/mem_cache/common.py | 3 +- .../model_runner_kv_cache_mixin.py | 9 +- .../srt/model_executor/pool_configurator.py | 109 ++++++++- python/sglang/srt/server_args.py | 8 + python/sglang/srt/utils/common.py | 21 ++ .../model_executor/test_pool_configurator.py | 231 +++++++++++++++++- 12 files changed, 398 insertions(+), 36 deletions(-) diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx index 89b7cfe70..a29020b97 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx @@ -5799,9 +5799,8 @@ do export GLOO_SOCKET_IFNAME=your_nic export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_NPU_FUSED_MOE_MODE=2 - export SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS=96 - python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ + python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} --disaggregation-decode-extra-slots 96 \ --cuda-graph-bs 8 16 24 32 40 \ --port 33000 --trust-remote-code \ --tp-size 16 --mem-fraction-static 0.76 --attention-backend ascend --device npu --quantization modelslim \ @@ -5934,9 +5933,8 @@ do export GLOO_SOCKET_IFNAME=your_nic export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_NPU_FUSED_MOE_MODE=2 - export SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS=96 - python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ + python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} --disaggregation-decode-extra-slots 96 \ --cuda-graph-bs 2 4 8 \ --port 33000 --trust-remote-code \ --tp-size 16 --mem-fraction-static 0.76 --attention-backend ascend --device npu --quantization modelslim \ diff --git a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx index a6915cd2f..574c6b00b 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx @@ -155,7 +155,6 @@ do export HCCL_BUFFSIZE=1600 export HCCL_SOCKET_IFNAME= export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=640 - export SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS=96 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_NPU_FUSED_MOE_MODE=2 @@ -172,6 +171,7 @@ do --device npu \ --quantization modelslim \ --disaggregation-transfer-backend ascend \ + --disaggregation-decode-extra-slots 96 \ --max-running-requests 80 \ --chunked-prefill-size -1 \ --moe-a2a-backend ascend_fuseep \ @@ -362,7 +362,6 @@ do export HCCL_BUFFSIZE=1600 export HCCL_SOCKET_IFNAME= export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=640 - export SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS=96 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_NPU_FUSED_MOE_MODE=2 @@ -379,6 +378,7 @@ do --device npu \ --quantization modelslim \ --disaggregation-transfer-backend ascend \ + --disaggregation-decode-extra-slots 96 \ --max-running-requests 80 \ --chunked-prefill-size -1 \ --moe-a2a-backend ascend_fuseep \ diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index 638654034..277bfe152 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -55,6 +55,20 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None: server_args.disable_radix_cache = True logger.warning("KV cache is forced as chunk cache for decode server") + # Default the number of *extra* decode req_to_token slots reserved for + # in-transfer (being-received-from-prefill) requests, on top of the + # max_running_requests-derived pool. Large batches get none; small + # per-worker batches reserve 2x the batch as cheap overlap headroom. + if server_args.disaggregation_decode_extra_slots is None: + extra_slots = 0 + if server_args.max_running_requests is not None: + per_worker = server_args.max_running_requests // max( + 1, server_args.dp_size + ) + if per_worker <= 32: + extra_slots = per_worker * 2 + server_args.disaggregation_decode_extra_slots = extra_slots + elif server_args.disaggregation_mode == "prefill": assert ( server_args.disaggregation_transfer_backend != "fake" diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 1b792e0ff..269da948d 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -378,9 +378,12 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): def _prealloc_required_tokens(self, req: Req) -> Tuple[int, int]: full_len, swa_len = self._prealloc_kv_lens(req) + swa_reserved = self.num_reserved_decode_tokens + if self.scheduler.server_args.disable_radix_cache: + swa_reserved = 0 return ( full_len + self.num_reserved_decode_tokens, - swa_len + self.num_reserved_decode_tokens, + swa_len + swa_reserved, ) def _init_kv_manager(self) -> CommonKVManager: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index eb2027bdf..44fc621c2 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -309,10 +309,6 @@ class Envs: SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS = EnvStr("{}") SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK = EnvBool(False) - # Extra slots in req_to_token_pool for decode workers (only effective when - # max_num_reqs > 32). Increases pool capacity so more KV cache transfers - # can overlap with decode execution without raising max_running_requests. - SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS = EnvInt(0) # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. @@ -338,7 +334,8 @@ class Envs: SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False) SGLANG_REQ_RUNNING_TIMEOUT = EnvFloat(-1) # in seconds SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120) - SGLANG_SWA_EVICTION_INTERVAL_MULTIPLIER = EnvFloat(1.0) + # Decode batches between SWA out-of-window evictions. + SGLANG_SWA_EVICTION_INTERVAL = EnvInt(128) # For non-streaming requests, the scheduler still flushes intermediate # output batches to the tokenizer manager every N decoded tokens so that # `first_token_time`/TTFT can be recorded. Lower this (e.g. to 1) to get diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 1da1a9684..e3b8c62f5 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2855,22 +2855,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): and hasattr(self.tree_cache, "dec_swa_lock_only") ) - # Eviction_interval: trade-off between SWA token waste and eviction overhead - page_size = self.tree_cache.page_size - eviction_interval = max( - page_size, - int( - sliding_window_size - * envs.SGLANG_SWA_EVICTION_INTERVAL_MULTIPLIER.get() - ), - ) - eviction_interval = (eviction_interval // page_size) * page_size + eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get()) + swa_maintenance_step = (self.forward_iter or 0) % eviction_interval == 0 for idx, req in enumerate(self.reqs): if self.forward_mode.is_decode(): # We set evict_swa condition here with two reasons: # 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running. - # 2. Evict swa every eviction_interval tokens to reduce the overhead. - if req.decode_batch_idx % eviction_interval == 1: + # 2. Evict swa every eviction_interval iterations to reduce the overhead. + if swa_maintenance_step and req.decode_batch_idx >= 1: self._evict_swa(req, req.seqlen - 1) # DSV4-NPU only (no-op elsewhere): the small paged compress-state @@ -2917,6 +2909,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): page_size=self.tree_cache.page_size, req_to_token_pool=self.req_to_token_pool, token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + drop_page_margin=self.tree_cache.is_chunk_cache(), ) def __str__(self): diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index e8a620b80..0d10b7947 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -71,6 +71,7 @@ def free_swa_out_of_window_slots( page_size: int, req_to_token_pool: ReqToTokenPool, token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, + drop_page_margin: bool = False, ) -> None: from sglang.srt.environ import envs @@ -86,7 +87,7 @@ def free_swa_out_of_window_slots( # preserving cache reuse in multi-turn scenarios. Without this, leaf nodes # may become tombstoned, causing SWA memory leak. # See also: _insert_helper case 3 in swa_radix_cache.py (defensive counterpart). - if envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get(): + if drop_page_margin or envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get(): evict_threshold = pre_len - sliding_window_size else: evict_threshold = pre_len - sliding_window_size - page_size diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index a10a54b88..685f1d405 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -328,13 +328,8 @@ class ModelRunnerKVCacheMixin: HybridMambaDecodeReqToTokenPool, ) - # subscribe memory for pre-allocated requests - # if max_num_reqs <= 32, we pre-allocate 2x requests - - pre_alloc_size = envs.SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS.get() - pre_alloc_size = ( - max_num_reqs * 2 if max_num_reqs <= 32 else pre_alloc_size - ) + # Extra slots for pre-allocated requests + pre_alloc_size = self.server_args.disaggregation_decode_extra_slots if config := self.mambaish_config: self.req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index fb1a7f095..dfa0c5b20 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -26,9 +26,14 @@ from sglang.srt.configs.model_config import ( ) from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.mem_cache.common import get_alloc_len_per_decode from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool -from sglang.srt.utils.common import is_float4_e2m1fn_x2 +from sglang.srt.utils.common import ( + ceil_align, + is_float4_e2m1fn_x2, + spec_decode_alloc_len_per_request, +) @dataclass @@ -331,6 +336,106 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): return self._solve_pool_sizes(max_total_num_tokens, page_size) +class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): + """Hybrid SWA configurator with the SWA pool sized from a fixed token cap. + + When max_running_requests is explicit, the SWA pool's worst-case + footprint is bounded per request. The SWA pool is sized tightly from that + cap and the freed memory is redirected to the full pool, instead of sizing + both pools by swa_full_tokens_ratio. + """ + + def __init__(self, mr: ModelRunner): + super().__init__(mr) + assert self._full_layers_num > 0 + + sa = mr.server_args + page_size = mr.page_size + window = mr.sliding_window_size + draft_tokens = sa.speculative_num_draft_tokens or 1 + eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get()) + + """ + __________[padding][eviction_interval][window] + Padding to make sure eviction point is page-aligned. + """ + trailing_tokens = window + eviction_interval * draft_tokens + page_size + if sa.speculative_algorithm is None: + decode_alloc = page_size + elif sa.disable_overlap_schedule: + # spec-v1: new_tokens_required_next_decode per request. + decode_alloc = spec_decode_alloc_len_per_request(sa) + else: + # spec-v2: the overlap allocator keeps 2 * alloc_len outstanding + # (eagle_info_v2.prepare_for_decode: kv_committed_len + 2 * alloc_len). + decode_alloc = 2 * get_alloc_len_per_decode(sa) + per_request = trailing_tokens + decode_alloc + + num_reqs = sa.max_running_requests // mr.dp_size + if sa.disaggregation_mode == "decode": + self._swa_cap = ( + per_request * num_reqs + + (window + page_size) * sa.disaggregation_decode_extra_slots + ) + else: + chunks_in_flight = 1 if sa.disable_overlap_schedule else 2 + self._swa_cap = ( + per_request * num_reqs + + chunks_in_flight * sa.chunked_prefill_size + + page_size + ) + + @staticmethod + def is_applicable(mr: ModelRunner) -> bool: + """True when SWAChunkCache can be sized from explicit max requests.""" + sa = mr.server_args + if sa.max_running_requests is None: + return False + if not sa.disable_radix_cache: + return False + if sa.chunked_prefill_size is None: + return False + if mr.sliding_window_size is None: + return False + return len(mr.model_config.full_attention_layer_ids) > 0 + + def calculate_pool_sizes( + self, available_bytes: int, page_size: int + ) -> MemoryPoolConfig: + # SWA pool sized tightly from the cap; the rest of the budget goes to full. + swa_tokens = ceil_align(self._swa_cap, page_size) + fixed_swa_bytes = swa_tokens * self._swa_per_token * self._swa_layers_num + full_cell_size = self._full_per_token * self._full_layers_num + full_tokens = ( + int((available_bytes - fixed_swa_bytes) // full_cell_size) // page_size + ) * page_size + if full_tokens <= 0: + raise RuntimeError( + f"SWA pool cap ({swa_tokens} tokens, " + f"{fixed_swa_bytes / (1 << 30):.2f} GiB) leaves no room for the full " + f"KV pool within the available {available_bytes / (1 << 30):.2f} GiB. " + f"Reduce --max-running-requests, lower SGLANG_SWA_EVICTION_INTERVAL, " + f"or increase --mem-fraction-static." + ) + return MemoryPoolConfig( + max_total_num_tokens=full_tokens, + full_max_total_num_tokens=full_tokens, + swa_max_total_num_tokens=swa_tokens, + ) + + def calculate_pool_sizes_from_max_tokens( + self, max_total_num_tokens: int, page_size: int + ) -> MemoryPoolConfig: + # Constrained max_total goes to the full pool; SWA stays at its cap. + swa_tokens = ceil_align(self._swa_cap, page_size) + full_tokens = (max_total_num_tokens // page_size) * page_size + return MemoryPoolConfig( + max_total_num_tokens=full_tokens, + full_max_total_num_tokens=full_tokens, + swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens), + ) + + @dataclass class _DSV4PoolSizes: full_max_total_num_tokens: int @@ -537,6 +642,8 @@ def create_memory_pool_configurator( if is_deepseek_v4(mr.model_config.hf_config) and mr.is_hybrid_swa: return DSV4PoolConfigurator(mr) if mr.is_hybrid_swa: + if SWAChunkCapPoolConfigurator.is_applicable(mr): + return SWAChunkCapPoolConfigurator(mr) return HybridSWAPoolConfigurator(mr) # Future: MambaPoolConfigurator return DefaultPoolConfigurator(mr) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 117b2b440..271c40c07 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1003,6 +1003,8 @@ class ServerArgs: disaggregation_decode_enable_radix_cache: bool = False disaggregation_decode_enable_offload_kvcache: bool = False num_reserved_decode_tokens: int = 512 # used for decode kv cache offload in PD + # Extra req_to_token slots for in-transfer requests; None -> default in PD hook + disaggregation_decode_extra_slots: Optional[int] = None # FIXME: hack to reduce ITL when decode bs is small disaggregation_decode_polling_interval: int = 1 optimistic_prefill_retries: int = 0 @@ -7348,6 +7350,12 @@ class ServerArgs: default=ServerArgs.num_reserved_decode_tokens, help="Number of decode tokens that will have memory reserved when adding new request to the running batch.", ) + parser.add_argument( + "--disaggregation-decode-extra-slots", + type=int, + default=ServerArgs.disaggregation_decode_extra_slots, + help="Number of extra decode req_to_token slots pre-allocated for in-transfer requests (PD mode). If unset, defaults to 0 (or 2x the per-worker running batch for small batches).", + ) parser.add_argument( "--disaggregation-decode-polling-interval", type=int, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 53f578252..9227c0a4a 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3416,6 +3416,27 @@ def ceil_align(x: int, y: int) -> int: return ceil_div(x, y) * y +def spec_decode_alloc_len_per_request(server_args) -> int: + """Per-request KV tokens a (spec-v1) decode step allocates: the draft-decode + topk*num_steps peak vs. the verify num_draft_tokens, page-aligned. + """ + page_size = server_args.page_size + len_per_topk = server_args.speculative_num_steps or 1 + spec_topk = server_args.speculative_eagle_topk or 1 + spec_tokens = server_args.speculative_num_draft_tokens or 1 + + if page_size > 1 and spec_topk > 1: + # last partial page and ceil alignment + len_per_topk = ceil_align(len_per_topk + page_size, page_size) + spec_tokens = ceil_align(spec_tokens, page_size) + elif page_size > 1: + # only page alignment + len_per_topk = ceil_align(len_per_topk, page_size) + spec_tokens = ceil_align(spec_tokens, page_size) + + return max(len_per_topk * spec_topk, spec_tokens) + + # COPIED FROM DeepGEMM def ceil_div(x: int, y: int) -> int: return (x + y - 1) // y diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 9d33f38ee..dbd9c24a4 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -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()