From 2ef00383abb8d111faf7a6746dd800de73b36aa2 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 8 Mar 2026 21:46:27 -0700 Subject: [PATCH] [Core] Refactor `init_memory_pool` into composable resolution helpers (#20142) --- .../sglang/srt/model_executor/model_runner.py | 10 -- .../model_runner_kv_cache_mixin.py | 140 ++++++++++-------- 2 files changed, 82 insertions(+), 68 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9b615cce8..ff3e44a6b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -602,16 +602,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Init memory pool and attention backends self.init_memory_pool(pre_model_load_memory) - # Init max running requests - self.max_running_requests = min( - ( - self.max_total_num_tokens // 2 - if server_args.max_running_requests is None - else server_args.max_running_requests // self.dp_size - ), - self.req_to_token_pool.size, - ) - # Init routed experts capturer self.init_routed_experts_capturer() 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 32cf5d7e8..a6b3f77a4 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 @@ -248,7 +248,9 @@ class ModelRunnerKVCacheMixin: return kv_cache_dim - def set_num_tokens_hybrid_swa(self: ModelRunner): + 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).""" page_size = self.server_args.page_size assert self.sliding_window_size is not None and self.sliding_window_size > 0 @@ -262,13 +264,12 @@ class ModelRunnerKVCacheMixin: if full_layers_num == 0: # all layers are SWA - self.swa_max_total_num_tokens = align_page_size(self.max_total_num_tokens) + self.swa_max_total_num_tokens = align_page_size(token_capacity) self.full_max_total_num_tokens = 0 - self.max_total_num_tokens = self.swa_max_total_num_tokens logger.info( f"Use sliding window memory pool (all SWA). swa_layer_tokens={self.swa_max_total_num_tokens}" ) - return + return self.swa_max_total_num_tokens swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio @@ -281,8 +282,8 @@ class ModelRunnerKVCacheMixin: # # The profile phase computed: # cell_size = F * n_full + S * n_swa - # max_total_num_tokens = rest_memory / cell_size - # => total_memory = max_total_num_tokens * (F * n_full + S * n_swa) + # token_capacity = rest_memory / cell_size + # => total_memory = token_capacity * (F * n_full + S * n_swa) # # We need to solve: # full_tokens * F * n_full + swa_tokens * S * n_swa = total_memory @@ -290,7 +291,7 @@ class ModelRunnerKVCacheMixin: # # Solution: # full_tokens = total_memory / (F * n_full + r * S * n_swa) - # = max_total_num_tokens * (F * n_full + S * n_swa) / (F * n_full + r * S * n_swa) + # = token_capacity * (F * n_full + S * n_swa) / (F * n_full + r * S * n_swa) kv_size = torch._utils._element_size(self.kv_cache_dtype) @@ -309,7 +310,7 @@ class ModelRunnerKVCacheMixin: ) # Total memory available from profile - total_memory = self.max_total_num_tokens * ( + total_memory = token_capacity * ( full_per_token * full_layers_num + swa_per_token * swa_layers_num ) @@ -322,19 +323,17 @@ 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 = int(total_memory / denominator) - self.swa_max_total_num_tokens = int( - self.full_max_total_num_tokens * 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) ) - - self.full_max_total_num_tokens = align_page_size(self.full_max_total_num_tokens) - self.swa_max_total_num_tokens = align_page_size(self.swa_max_total_num_tokens) - - self.max_total_num_tokens = self.full_max_total_num_tokens 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}" ) + return self.full_max_total_num_tokens def _calculate_mamba_ratio(self: ModelRunner) -> int: if self.server_args.disable_radix_cache: @@ -676,64 +675,89 @@ class ModelRunnerKVCacheMixin: self.token_to_kv_pool_allocator.full_to_swa_index_mapping ) - def init_memory_pool(self: ModelRunner, pre_model_load_memory: int): - max_total_tokens = self.server_args.max_total_tokens - self.max_total_num_tokens = self.profile_max_num_token(pre_model_load_memory) + def _resolve_token_capacity(self: ModelRunner, profiled_tokens: int) -> int: + """Compute final token pool capacity from profiled value, + applying user cap, page alignment, and PP sync""" + user_limit = self.server_args.max_total_tokens - # Resolve max_num_reqs (per dp worker) - max_num_reqs = self.server_args.max_running_requests - if max_num_reqs is not None: - max_num_reqs = max_num_reqs // self.dp_size - else: - estimated = int( - self.max_total_num_tokens / self.model_config.context_len * 512 - ) - max_num_reqs = max(min(estimated, 4096), 2048) - - if self.mambaish_config is not None: - ratio = self._calculate_mamba_ratio() - - # Constrain the max_num_reqs by the mamba cache size - max_num_reqs = min( - max_num_reqs, self.server_args.max_mamba_cache_size // ratio - ) - - if max_total_tokens is not None: - if max_total_tokens > self.max_total_num_tokens: + # Apply user-specified upper bound + if user_limit is not None: + if user_limit > profiled_tokens: logging.warning( - f"max_total_tokens={max_total_tokens} is larger than the profiled value " - f"{self.max_total_num_tokens}. " - f"Use the profiled value instead." + f"max_total_tokens={user_limit} is larger than the profiled value " + f"{profiled_tokens}. Use the profiled value instead." ) - self.max_total_num_tokens = min(self.max_total_num_tokens, max_total_tokens) + capacity = min(profiled_tokens, user_limit) + else: + capacity = profiled_tokens - self.max_total_num_tokens = ( - self.max_total_num_tokens - // self.server_args.page_size - * self.server_args.page_size - ) - # different pp rank may have different num of layers, so we need to reduce the max_total_num_tokens + # Align to page boundary + page_size = self.server_args.page_size + capacity = capacity // page_size * page_size + + # Sync across PP ranks (each may have different layer counts) if self.pp_size > 1: - tensor = torch.tensor(self.max_total_num_tokens, dtype=torch.int64) + tensor = torch.tensor(capacity, dtype=torch.int64) torch.distributed.all_reduce( tensor, op=torch.distributed.ReduceOp.MIN, group=get_world_group().cpu_group, ) - self.max_total_num_tokens = tensor.item() + capacity = tensor.item() + return capacity + + def _resolve_max_num_reqs(self: ModelRunner, token_capacity: int) -> int: + """Compute max concurrent requests (per dp worker) from the finalized + token capacity.""" + # Estimate pool size (used as upper bound when user specifies max_running_requests) + estimated = int(token_capacity / self.model_config.context_len * 512) + estimated = max(min(estimated, 4096), 2048) + + max_num_reqs = self.server_args.max_running_requests + if max_num_reqs is not None: + max_num_reqs = min(max_num_reqs // self.dp_size, estimated) + else: + max_num_reqs = min(estimated, token_capacity // 2) + + if self.mambaish_config is not None: + ratio = self._calculate_mamba_ratio() + max_num_reqs = min( + max_num_reqs, self.server_args.max_mamba_cache_size // ratio + ) + + 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) + + # Resolve the token capacity + 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: - self.max_total_num_tokens = self.server_args.draft_runner_cache_size - max_num_reqs = self.server_args.max_num_reqs + token_capacity = self.server_args.draft_runner_cache_size - # create token size for hybrid cache + # Hybrid SWA: split capacity into full/swa pools, adjust effective capacity if self.is_hybrid_swa: - self.set_num_tokens_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: - # Draft worker should use SWA adjusted max_total_num_tokens for cache size, otherwise it may cause oob in kv cache store self.server_args.draft_runner_cache_size = self.max_total_num_tokens - self.server_args.max_num_reqs = max_num_reqs + self.server_args.max_num_reqs = self.max_running_requests if self.max_total_num_tokens <= 0: raise RuntimeError( @@ -741,7 +765,7 @@ class ModelRunnerKVCacheMixin: f"Current value: {self.server_args.mem_fraction_static=}" ) - self._init_pools(max_num_reqs) + self._init_pools(self.max_running_requests) logger.info( f"Memory pool end. "