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

Adjust mamba cache allocation (#18786)

Browse Source
This commit is contained in:
Ke Bao
2026-02-13 18:06:23 +08:00
committed by GitHub
parent a6c4b52ac5
commit eacab2868a
1 changed files with 11 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -170,20 +170,21 @@ class ModelRunnerKVCacheMixin:
mamba_state_intermediate_size / (1 << 30)
)
if (
server_args.disable_radix_cache
or server_args.max_mamba_cache_size is not None
):
# with disable radix cache, sets the max_mamba_cache_size based on the max_running_requests
if server_args.max_mamba_cache_size is None:
if server_args.max_running_requests is not None:
server_args.max_mamba_cache_size = server_args.max_running_requests
else:
server_args.max_mamba_cache_size = 512
if server_args.max_mamba_cache_size is not None:
# Use explicitly set max_mamba_cache_size
server_args.max_mamba_cache_size = server_args.max_mamba_cache_size // (
server_args.dp_size if server_args.enable_dp_attention else 1
)
elif (
server_args.disable_radix_cache
and server_args.max_running_requests is not None
):
# Use explicitly set max_running_requests when radix cache is disabled
server_args.max_mamba_cache_size = server_args.max_running_requests // (
server_args.dp_size if server_args.enable_dp_attention else 1
)
else:
# Use ratio-based calculation to auto-fit available memory
assert config.mamba2_cache_params.mamba_cache_per_req > 0
# allocate the memory based on the ratio between mamba state memory vs. full kv cache memory