chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Core] Extract _calculate_mamba_ratio and _init_pools from init_memory_pool (#20058)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-03-06 13:37:22 -08:00
committed by GitHub
parent 8cdb7e1fd4
commit 7a6cf0e9ba
2 changed files with 89 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/model_executor/model_runner.py
View File

@@ -607,7 +607,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
(
self.max_total_num_tokens // 2
if server_args.max_running_requests is None
else server_args.max_running_requests // (self.dp_size)
else server_args.max_running_requests // self.dp_size
),
self.req_to_token_pool.size,
)

166
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -336,88 +336,20 @@ class ModelRunnerKVCacheMixin:
f"Use sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}"
)
def init_memory_pool(self: ModelRunner, total_gpu_memory: int):
max_num_reqs = self.server_args.max_running_requests
max_total_tokens = self.server_args.max_total_tokens
self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
def _calculate_mamba_ratio(self: ModelRunner) -> int:
if self.server_args.disable_radix_cache:
return 1
if max_num_reqs is None:
max_num_reqs = min(
max(
int(
self.max_total_num_tokens / self.model_config.context_len * 512
),
2048,
),
4096,
)
if self.mambaish_config is not None:
additional_ratio = 0
if self.server_args.enable_mamba_extra_buffer():
if not self.spec_algorithm.is_none():
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP
else:
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP
if self.server_args.disable_radix_cache:
ratio = 1
additional_ratio = 0
if self.server_args.enable_mamba_extra_buffer():
if not self.spec_algorithm.is_none():
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP
else:
ratio = MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
max_num_reqs = min(
max_num_reqs, self.server_args.max_mamba_cache_size // ratio
)
# for dp attention, we need control the max_num_reqs for speculative decoding mamba space
if (
not self.spec_algorithm.is_none()
and self.server_args.enable_dp_attention
):
max_num_reqs = min(
max_num_reqs, self.server_args.max_running_requests // self.dp_size
)
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP
if max_total_tokens is not None:
if max_total_tokens > self.max_total_num_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."
)
self.max_total_num_tokens = min(self.max_total_num_tokens, max_total_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
if self.pp_size > 1:
tensor = torch.tensor(self.max_total_num_tokens, 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()
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
# create token size for hybrid cache
if self.is_hybrid_swa:
self.set_num_tokens_hybrid_swa()
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
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=}"
)
return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
def _init_pools(self: ModelRunner, max_num_reqs: int):
# 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
@@ -743,6 +675,84 @@ class ModelRunnerKVCacheMixin:
self.token_to_kv_pool_allocator.full_to_swa_index_mapping
)
def init_memory_pool(self: ModelRunner, total_gpu_memory: int):
max_num_reqs = self.server_args.max_running_requests
max_total_tokens = self.server_args.max_total_tokens
self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
if max_num_reqs is None:
max_num_reqs = min(
max(
int(
self.max_total_num_tokens / self.model_config.context_len * 512
),
2048,
),
4096,
)
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
)
# for dp attention, we need control the max_num_reqs for speculative decoding mamba space
if (
not self.spec_algorithm.is_none()
and self.server_args.enable_dp_attention
):
max_num_reqs = min(
max_num_reqs, self.server_args.max_running_requests // self.dp_size
)
if max_total_tokens is not None:
if max_total_tokens > self.max_total_num_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."
)
self.max_total_num_tokens = min(self.max_total_num_tokens, max_total_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
if self.pp_size > 1:
tensor = torch.tensor(self.max_total_num_tokens, 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()
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
# create token size for hybrid cache
if self.is_hybrid_swa:
self.set_num_tokens_hybrid_swa()
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
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(max_num_reqs)
logger.info(
f"Memory pool end. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"