[Core] Replace server_args mutation hack with explicit MemoryPoolConfig for draft worker init (#20183)

python/sglang/srt/managers/tp_worker.py

@@ -54,6 +54,7 @@ from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
 if TYPE_CHECKING:
     from sglang.srt.managers.cache_controller import LayerDoneCounter
     from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.model_runner_kv_cache_mixin import MemoryPoolConfig
 
 logger = logging.getLogger(__name__)
 
@@ -231,6 +232,7 @@ class TpModelWorker(BaseTpWorker):
         is_draft_worker: bool = False,
         req_to_token_pool: Optional[ReqToTokenPool] = None,
         token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
+        memory_pool_config: Optional[MemoryPoolConfig] = None,
         is_multi_layer_eagle: bool = False,
     ):
         # Parse args
@@ -248,6 +250,7 @@ class TpModelWorker(BaseTpWorker):
         self.is_multi_layer_eagle = is_multi_layer_eagle
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.memory_pool_config = memory_pool_config
         self.attn_cp_rank = attn_cp_rank
         self.moe_dp_rank = moe_dp_rank
 
@@ -354,6 +357,7 @@ class TpModelWorker(BaseTpWorker):
             is_draft_worker=self.is_draft_worker,
             req_to_token_pool=self.req_to_token_pool,
             token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            memory_pool_config=self.memory_pool_config,
             draft_model_idx=0 if self.is_multi_layer_eagle else None,
         )
 
@@ -379,6 +383,7 @@ class TpModelWorker(BaseTpWorker):
                     is_draft_worker=self.is_draft_worker,
                     req_to_token_pool=self.req_to_token_pool,
                     token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                    memory_pool_config=self.memory_pool_config,
                     draft_model_idx=i,
                 )
             )

python/sglang/srt/model_executor/model_runner.py

@@ -13,6 +13,8 @@
 # ==============================================================================
 """ModelRunner runs the forward passes of the models."""
 
+from __future__ import annotations
+
 import datetime
 import gc
 import inspect
@@ -132,6 +134,7 @@ from sglang.srt.model_executor.forward_batch_info import (
 )
 from sglang.srt.model_executor.hook_manager import register_forward_hooks
 from sglang.srt.model_executor.model_runner_kv_cache_mixin import (
+    MemoryPoolConfig,
     ModelRunnerKVCacheMixin,
 )
 from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
@@ -301,6 +304,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         is_draft_worker: bool = False,
         req_to_token_pool: Optional[ReqToTokenPool] = None,
         token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
+        memory_pool_config: Optional[MemoryPoolConfig] = None,
         draft_model_idx: Optional[int] = None,
     ):
         # Parse args
@@ -322,6 +326,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self.dist_port = nccl_port
         self.server_args = server_args
         self.is_draft_worker = is_draft_worker
+        self.memory_pool_config = memory_pool_config
         self.is_generation = model_config.is_generation
         self.is_multimodal = model_config.is_multimodal
         self.is_multimodal_chunked_prefill_supported = (

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
 
@@ -33,6 +34,26 @@ from sglang.srt.utils.common import (
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
 
+
+@dataclass
+class MemoryPoolConfig:
+    """Resolved memory pool config, shared between target and draft workers."""
+
+    max_total_num_tokens: int
+    max_running_requests: int
+    full_max_total_num_tokens: Optional[int] = None
+    swa_max_total_num_tokens: Optional[int] = None
+
+    mem_fraction_static: Optional[float] = None
+
+    def __post_init__(self):
+        if self.max_total_num_tokens <= 0:
+            msg = "Not enough memory. Please try to increase --mem-fraction-static."
+            if self.mem_fraction_static is not None:
+                msg += f" Current value: mem_fraction_static={self.mem_fraction_static}"
+            raise RuntimeError(msg)
+
+
 # the ratio of mamba cache pool size to max_running_requests
 MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
 MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2
@@ -248,9 +269,13 @@ class ModelRunnerKVCacheMixin:
 
         return kv_cache_dim
 
-    def set_num_tokens_hybrid_swa(self: ModelRunner, token_capacity: int) -> int:
+    def _resolve_hybrid_swa_tokens(
-        """Split token_capacity into full/swa pools. Returns the effective
+        self: ModelRunner, token_capacity: int
-        max_total_num_tokens (= full pool size)."""
+    ) -> Tuple[int, int, int]:
+        """Split token_capacity into full/swa pools.
+
+        Returns (effective_capacity, full_max_total_num_tokens, swa_max_total_num_tokens).
+        """
         page_size = self.server_args.page_size
 
         assert self.sliding_window_size is not None and self.sliding_window_size > 0
@@ -264,12 +289,11 @@ class ModelRunnerKVCacheMixin:
 
         if full_layers_num == 0:
             # all layers are SWA
-            self.swa_max_total_num_tokens = align_page_size(token_capacity)
+            swa_tokens = align_page_size(token_capacity)
-            self.full_max_total_num_tokens = 0
             logger.info(
-                f"Use sliding window memory pool (all SWA). swa_layer_tokens={self.swa_max_total_num_tokens}"
+                f"Use sliding window memory pool (all SWA). swa_layer_tokens={swa_tokens}"
             )
-            return self.swa_max_total_num_tokens
+            return swa_tokens, 0, swa_tokens
 
         swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio
 
@@ -323,17 +347,13 @@ class ModelRunnerKVCacheMixin:
             denominator > 0
         ), f"Invalid denominator={denominator} for memory-based allocation. full_per_token={full_per_token}, full_layers_num={full_layers_num}, swa_per_token={swa_per_token}, swa_layers_num={swa_layers_num}, swa_full_tokens_ratio={swa_full_tokens_ratio}"
 
-        self.full_max_total_num_tokens = align_page_size(
+        full_tokens = align_page_size(int(total_memory / denominator))
-            int(total_memory / denominator)
+        swa_tokens = align_page_size(int(full_tokens * swa_full_tokens_ratio))
-        )
-        self.swa_max_total_num_tokens = align_page_size(
-            int(self.full_max_total_num_tokens * swa_full_tokens_ratio)
-        )
 
         logger.info(
-            f"Use sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}"
+            f"Use sliding window memory pool. full_layer_tokens={full_tokens}, swa_layer_tokens={swa_tokens}"
         )
-        return self.full_max_total_num_tokens
+        return full_tokens, full_tokens, swa_tokens
 
     def _calculate_mamba_ratio(self: ModelRunner) -> int:
         if self.server_args.disable_radix_cache:
@@ -349,7 +369,10 @@ class ModelRunnerKVCacheMixin:
 
         return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
 
-    def _init_pools(self: ModelRunner, max_num_reqs: int):
+    def _init_pools(self: ModelRunner):
+        """Initialize the memory pools."""
+        max_num_reqs = self.max_running_requests
+
         # Initialize req_to_token_pool
         if self.req_to_token_pool is None:
             # FIXME(lsyin): this is the temporary fix for the context length issue when using speculative decoding
@@ -728,44 +751,49 @@ class ModelRunnerKVCacheMixin:
 
         return max_num_reqs
 
-    def init_memory_pool(self: ModelRunner, pre_model_load_memory: int):
+    def _apply_memory_pool_config(self: ModelRunner, config: MemoryPoolConfig):
-        # Profile the maximum number of tokens
+        """Apply a resolved MemoryPoolConfig and initialize pools."""
-        profiled_tokens = self.profile_max_num_token(pre_model_load_memory)
+        self.max_total_num_tokens = config.max_total_num_tokens
+        self.max_running_requests = config.max_running_requests
+        if self.is_hybrid_swa:
+            self.full_max_total_num_tokens = config.full_max_total_num_tokens
+            self.swa_max_total_num_tokens = config.swa_max_total_num_tokens
 
-        # Resolve the token capacity
+        self._init_pools()
+
+    def _resolve_memory_pool_config(
+        self: ModelRunner, pre_model_load_memory: int
+    ) -> MemoryPoolConfig:
+        """Profile GPU memory and resolve all pool parameters into a config."""
+        profiled_tokens = self.profile_max_num_token(pre_model_load_memory)
         token_capacity = self._resolve_token_capacity(profiled_tokens)
 
-        # HACK: spec decode uses server_args as a mutable channel to pass
+        full_tokens = None
-        # resolved values between target and draft workers. Target writes first,
+        swa_tokens = None
-        # draft reads later. Should be replaced with an explicit handoff.
-        # NOTE: draft worker override must happen BEFORE SWA splitting so that
-        # swa_max_total_num_tokens is computed from the correct base value.
-        if not self.spec_algorithm.is_none() and self.is_draft_worker:
-            token_capacity = self.server_args.draft_runner_cache_size
-
-        # Hybrid SWA: split capacity into full/swa pools, adjust effective capacity
         if self.is_hybrid_swa:
-            token_capacity = self.set_num_tokens_hybrid_swa(token_capacity)
+            token_capacity, full_tokens, swa_tokens = self._resolve_hybrid_swa_tokens(
-
+                token_capacity
-        # Commit the resolved token capacity & max number of requests
-        self.max_total_num_tokens = token_capacity
-        if not self.spec_algorithm.is_none() and self.is_draft_worker:
-            self.max_running_requests = self.server_args.max_num_reqs
-        else:
-            self.max_running_requests = self._resolve_max_num_reqs(token_capacity)
-
-        # Target worker stores resolved values for draft worker to read later
-        if not self.spec_algorithm.is_none() and not self.is_draft_worker:
-            self.server_args.draft_runner_cache_size = self.max_total_num_tokens
-            self.server_args.max_num_reqs = self.max_running_requests
-
-        if self.max_total_num_tokens <= 0:
-            raise RuntimeError(
-                f"Not enough memory. Please try to increase --mem-fraction-static. "
-                f"Current value: {self.server_args.mem_fraction_static=}"
             )
 
-        self._init_pools(self.max_running_requests)
+        return MemoryPoolConfig(
+            max_total_num_tokens=token_capacity,
+            max_running_requests=self._resolve_max_num_reqs(token_capacity),
+            full_max_total_num_tokens=full_tokens,
+            swa_max_total_num_tokens=swa_tokens,
+            mem_fraction_static=self.server_args.mem_fraction_static,
+        )
+
+    def init_memory_pool(self: ModelRunner, pre_model_load_memory: int):
+        if not self.spec_algorithm.is_none() and self.is_draft_worker:
+            assert (
+                self.memory_pool_config is not None
+            ), "Draft worker requires memory_pool_config"
+        else:
+            self.memory_pool_config = self._resolve_memory_pool_config(
+                pre_model_load_memory
+            )
+
+        self._apply_memory_pool_config(self.memory_pool_config)
 
         logger.info(
             f"Memory pool end. "

python/sglang/srt/speculative/eagle_worker.py

@@ -152,6 +152,7 @@ class EAGLEWorker(TpModelWorker):
                 is_draft_worker=True,
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                memory_pool_config=target_worker.model_runner.memory_pool_config,
             )
 
         embed, head = self.target_worker.model_runner.model.get_embed_and_head()

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -145,6 +145,7 @@ class EagleDraftWorker(BaseDraftWorker):
                 is_draft_worker=True,
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                memory_pool_config=target_worker.model_runner.memory_pool_config,
             )
 
         # Alias for better readability

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -142,6 +142,7 @@ class MultiLayerEagleWorker(TpModelWorker):
                 is_draft_worker=True,
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                memory_pool_config=target_worker.model_runner.memory_pool_config,
                 is_multi_layer_eagle=True,
             )
 

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -126,6 +126,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
                 is_draft_worker=True,
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                memory_pool_config=target_worker.model_runner.memory_pool_config,
                 is_multi_layer_eagle=True,
             )
 

python/sglang/srt/speculative/standalone_worker.py

@@ -86,6 +86,7 @@ class StandaloneWorker(EAGLEWorker):
                 is_draft_worker=True,
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                memory_pool_config=target_worker.model_runner.memory_pool_config,
             )
 
         # Init attention backend and cuda graphs

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -99,6 +99,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
                 is_draft_worker=True,
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                memory_pool_config=target_worker.model_runner.memory_pool_config,
             )
 
         # Alias for better readability