diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 8408cc905..8f5500639 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/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, ) ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d6cf44b43..e42dbc556 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 = ( 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 a6b3f77a4..577508b1f 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 @@ -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: - """Split token_capacity into full/swa pools. Returns the effective - max_total_num_tokens (= full pool size).""" + def _resolve_hybrid_swa_tokens( + self: ModelRunner, token_capacity: int + ) -> 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) - self.full_max_total_num_tokens = 0 + swa_tokens = align_page_size(token_capacity) 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( - int(total_memory / denominator) - ) - self.swa_max_total_num_tokens = align_page_size( - int(self.full_max_total_num_tokens * swa_full_tokens_ratio) - ) + full_tokens = align_page_size(int(total_memory / denominator)) + swa_tokens = align_page_size(int(full_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): - # Profile the maximum number of tokens - profiled_tokens = self.profile_max_num_token(pre_model_load_memory) + def _apply_memory_pool_config(self: ModelRunner, config: MemoryPoolConfig): + """Apply a resolved MemoryPoolConfig and initialize pools.""" + 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 - # resolved values between target and draft workers. Target writes first, - # 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 + full_tokens = None + swa_tokens = None if self.is_hybrid_swa: - token_capacity = self.set_num_tokens_hybrid_swa(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=}" + token_capacity, full_tokens, swa_tokens = self._resolve_hybrid_swa_tokens( + token_capacity ) - 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. " diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index b5277fb88..6e1a41249 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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() diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 5aa4d2c63..c32d8a83f 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index c5ea19981..e9ee57f89 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 32f264361..fbad4adb9 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/standalone_worker.py b/python/sglang/srt/speculative/standalone_worker.py index da21b989b..b61e67db4 100644 --- a/python/sglang/srt/speculative/standalone_worker.py +++ b/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 diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index f5ee84508..ecaab3cc0 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/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