Adjust mamba cache allocation (#18786)
This commit is contained in:
@@ -170,20 +170,21 @@ class ModelRunnerKVCacheMixin:
|
|||||||
mamba_state_intermediate_size / (1 << 30)
|
mamba_state_intermediate_size / (1 << 30)
|
||||||
)
|
)
|
||||||
|
|
||||||
if (
|
if server_args.max_mamba_cache_size is not None:
|
||||||
server_args.disable_radix_cache
|
# Use explicitly set max_mamba_cache_size
|
||||||
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
|
|
||||||
server_args.max_mamba_cache_size = server_args.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
|
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:
|
else:
|
||||||
|
# Use ratio-based calculation to auto-fit available memory
|
||||||
assert config.mamba2_cache_params.mamba_cache_per_req > 0
|
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
|
# allocate the memory based on the ratio between mamba state memory vs. full kv cache memory
|
||||||
|
|||||||
Reference in New Issue
Block a user