From 6a62eabed626ef66f20c26978f2481414dd1bc0d Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 4 May 2026 16:04:31 -0700 Subject: [PATCH] consolidate NSA pool construction (#24389) Co-authored-by: Zhiqiang Xie --- .../model_runner_kv_cache_mixin.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) 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 58658b12c..9bc7e123c 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 @@ -412,8 +412,18 @@ class ModelRunnerKVCacheMixin: end_layer=self.end_layer, ) elif self.use_mla_backend and is_nsa_model: - nsa_pool_kwargs = dict( - size=self.max_total_num_tokens, + PoolCls = ( + HiSparseNSATokenToKVPool if self.enable_hisparse else NSATokenToKVPool + ) + pool_kwargs = {} + if self.enable_hisparse: + from sglang.srt.mem_cache.sparsity import parse_hisparse_config + + pool_kwargs["host_to_device_ratio"] = parse_hisparse_config( + self.server_args + ).host_to_device_ratio + self.token_to_kv_pool = PoolCls( + self.max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, @@ -425,17 +435,8 @@ class ModelRunnerKVCacheMixin: start_layer=self.start_layer, end_layer=self.end_layer, index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config), + **pool_kwargs, ) - if self.enable_hisparse: - from sglang.srt.mem_cache.sparsity import parse_hisparse_config - - hisparse_cfg = parse_hisparse_config(self.server_args) - nsa_pool_kwargs["host_to_device_ratio"] = ( - hisparse_cfg.host_to_device_ratio - ) - self.token_to_kv_pool = HiSparseNSATokenToKVPool(**nsa_pool_kwargs) - else: - self.token_to_kv_pool = NSATokenToKVPool(**nsa_pool_kwargs) elif self.use_mla_backend and not self.mambaish_config: assert not is_nsa_model if is_float4_e2m1fn_x2(self.kv_cache_dtype):