Cap SWA pool sizing with chunk cache (#28755)

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"

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:

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

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):

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

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,

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)

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,

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