consolidate NSA pool construction (#24389)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

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(
+            PoolCls = (
-                size=self.max_total_num_tokens,
+                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):