From 4e35c30cbe4757cd691aaea56c964adddb05f68f Mon Sep 17 00:00:00 2001 From: Taegeon Um Date: Wed, 13 May 2026 12:45:08 +0900 Subject: [PATCH] [Bugfix, NSA HiCache] Fix missing override_kv_cache_dim in attach_hybrid_nsa_pool_to_hiradix_cache (#25022) Co-authored-by: Taegeon Um Co-authored-by: Zhangheng --- .../hybrid_cache/hybrid_cache_controller.py | 2 -- .../hybrid_cache/hybrid_pool_assembler.py | 32 ++++++------------- .../srt/mem_cache/unified_radix_cache.py | 2 ++ 3 files changed, 12 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 456bd72c8..b4fdde959 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -164,8 +164,6 @@ class HybridCacheController(BaseHiCacheController): storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, - attn_cp_rank: int = 0, - attn_cp_size: int = 1, transfer_layer_num: Optional[int] = None, enable_storage_metrics: bool = False, ): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index c4d857b5d..0a8b10e42 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -110,8 +110,6 @@ def build_kv_only_stack( storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, - attn_cp_rank: int = 0, - attn_cp_size: int = 1, enable_storage_metrics: bool = False, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping) @@ -149,8 +147,6 @@ def build_kv_only_stack( storage_backend_extra_config=storage_backend_extra_config, pp_rank=pp_rank, pp_size=pp_size, - attn_cp_rank=attn_cp_rank, - attn_cp_size=attn_cp_size, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, ) @@ -168,6 +164,8 @@ def build_hybrid_swa_stack( page_size: int, tp_group, load_cache_event, + attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, storage_backend: Optional[str], use_mla: bool, host_swa_evict_fn: Optional[Callable[[int], Any]] = None, @@ -177,8 +175,6 @@ def build_hybrid_swa_stack( storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, - attn_cp_rank: int = 0, - attn_cp_size: int = 1, enable_storage_metrics: bool = False, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) @@ -225,6 +221,8 @@ def build_hybrid_swa_stack( page_size, tp_group, load_cache_event=load_cache_event, + attn_cp_group=attn_cp_group, + attn_tp_group=attn_tp_group, write_policy=server_args.hicache_write_policy, io_backend=server_args.hicache_io_backend, storage_backend=storage_backend, @@ -233,8 +231,6 @@ def build_hybrid_swa_stack( storage_backend_extra_config=storage_backend_extra_config, pp_rank=pp_rank, pp_size=pp_size, - attn_cp_rank=attn_cp_rank, - attn_cp_size=attn_cp_size, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, ) @@ -263,8 +259,6 @@ def build_hybrid_mamba_stack( storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, - attn_cp_rank: int = 0, - attn_cp_size: int = 1, enable_storage_metrics: bool = False, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping) @@ -317,8 +311,6 @@ def build_hybrid_mamba_stack( storage_backend_extra_config=storage_backend_extra_config, pp_rank=pp_rank, pp_size=pp_size, - attn_cp_rank=attn_cp_rank, - attn_cp_size=attn_cp_size, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, ) @@ -346,8 +338,6 @@ def build_shared_anchor_stack( storage_backend_extra_config: Optional[dict] = None, pp_rank: int = 0, pp_size: int = 1, - attn_cp_rank: int = 0, - attn_cp_size: int = 1, enable_storage_metrics: bool = False, ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping) @@ -394,8 +384,6 @@ def build_shared_anchor_stack( storage_backend_extra_config=storage_backend_extra_config, pp_rank=pp_rank, pp_size=pp_size, - attn_cp_rank=attn_cp_rank, - attn_cp_size=attn_cp_size, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, ) @@ -505,6 +493,8 @@ def attach_hybrid_pool_to_unified_cache( page_size=cache.page_size, tp_group=params.tp_cache_group, load_cache_event=load_cache_event, + attn_cp_group=attn_cp_group, + attn_tp_group=attn_tp_group, storage_backend=None, use_mla=False, host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA), @@ -538,8 +528,11 @@ def attach_hybrid_pool_to_unified_cache( page_size=cache.page_size, tp_group=params.tp_cache_group, load_cache_event=load_cache_event, + attn_cp_group=attn_cp_group, + attn_tp_group=attn_tp_group, storage_backend=None, use_mla=use_mla, + override_kv_cache_dim=full_kv_pool.kv_cache_dim, shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( full_kv_pool, kv_host_pool, @@ -548,8 +541,6 @@ def attach_hybrid_pool_to_unified_cache( ), pp_rank=params.pp_rank, pp_size=params.pp_size, - attn_cp_rank=params.attn_cp_rank, - attn_cp_size=params.attn_cp_size, ) cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) cache.host_pool_group = host_pool_group @@ -645,6 +636,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( attn_tp_group=attn_tp_group, storage_backend=server_args.hicache_storage_backend, use_mla=True, + override_kv_cache_dim=kv.kv_cache_dim, prefetch_threshold=prefetch_threshold, shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( kv, @@ -656,8 +648,6 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( storage_backend_extra_config=extra_config, pp_rank=radix_cache.pp_rank, pp_size=radix_cache.pp_size, - attn_cp_rank=params.attn_cp_rank, - attn_cp_size=params.attn_cp_size, enable_storage_metrics=enable_storage_metrics, ) radix_cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) @@ -715,8 +705,6 @@ def attach_hybrid_pool_to_mamba_cache( storage_backend_extra_config=extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, - attn_cp_rank=params.attn_cp_rank, - attn_cp_size=params.attn_cp_size, enable_storage_metrics=enable_storage_metrics, ) mamba_cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 50fdd2ef6..c1618f6a4 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -319,6 +319,8 @@ class UnifiedRadixCache(BasePrefixCache): params, server_args, load_cache_event=self.load_cache_event, + attn_cp_group=params.attn_cp_cache_group, + attn_tp_group=params.attn_tp_cache_group, ) # State initialization