[HiCache & HybridModel] mooncake backend support DSA & mamba model

python/sglang/srt/mem_cache/hicache_storage.py

@@ -39,12 +39,13 @@ class PoolName(str, Enum):
 
     KV = "kv"
     MAMBA = "mamba"
+    INDEXER = "indexer"
 
 
 class PoolHitPolicy(str, Enum):
     """Hit policy for batch_exists_v2 per-pool prefix matching.
 
-    ALL_PAGES      : every page in [0, kv_hit) must exist (default).
+    ALL_PAGES      : every page in [0, kv_hit) must exist (e.g. DSA).
     TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states).
     """
 

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -25,6 +25,17 @@
     MatchPrefixParams,
     MatchResult,
 )
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
+    PoolName,
+    PoolTransfer,
+)
+from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
+    HybridCacheController,
+)
+from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
+    build_nsa_hybrid_stack,
+)
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
     MLATokenToKVPool,
@@ -33,7 +44,6 @@
 from sglang.srt.mem_cache.memory_pool_host import (
     MHATokenToKVPoolHost,
     MLATokenToKVPoolHost,
-    NSATokenToKVPoolHost,
 )
 from sglang.srt.mem_cache.radix_cache import (
     RadixCache,
@@ -70,14 +80,8 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs):
                 allocator_type=server_args.hicache_storage_backend,
             )
         elif isinstance(self.kv_cache, NSATokenToKVPool):
-            self.token_to_kv_pool_host = NSATokenToKVPoolHost(
-                self.kv_cache,
-                server_args.hicache_ratio,
-                server_args.hicache_size,
-                self.page_size,
-                server_args.hicache_mem_layout,
-                allocator_type=server_args.hicache_storage_backend,
-            )
+            # Filled by build_nsa_hybrid_stack after storage extra_config is parsed.
+            self.token_to_kv_pool_host = None
         elif isinstance(self.kv_cache, MLATokenToKVPool):
             self.token_to_kv_pool_host = MLATokenToKVPoolHost(
                 self.kv_cache,
@@ -88,7 +92,9 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs):
                 allocator_type=server_args.hicache_storage_backend,
             )
         else:
-            raise ValueError(f"HiRadixCache only supports MHA and MLA yet")
+            raise ValueError(
+                "HiRadixCache only supports MHA, MLA, and NSA (DSA) models"
+            )
 
         self.tp_group = params.tp_cache_group
         self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
@@ -112,22 +118,33 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs):
         self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
 
         self.load_cache_event = threading.Event()
-        self.cache_controller = HiCacheController(
-            params.token_to_kv_pool_allocator,
-            self.token_to_kv_pool_host,
-            self.page_size,
-            self.tp_group,
-            load_cache_event=self.load_cache_event,
-            write_policy=server_args.hicache_write_policy,
-            io_backend=server_args.hicache_io_backend,
-            storage_backend=server_args.hicache_storage_backend,
-            prefetch_threshold=prefetch_threshold,
-            model_name=server_args.served_model_name,
-            storage_backend_extra_config=extra_config,
-            pp_rank=self.pp_rank,
-            pp_size=self.pp_size,
-            enable_storage_metrics=self.enable_storage_metrics,
-        )
+        if isinstance(self.kv_cache, NSATokenToKVPool):
+            build_nsa_hybrid_stack(
+                self,
+                params,
+                server_args,
+                extra_config=extra_config,
+                prefetch_threshold=prefetch_threshold,
+                enable_storage_metrics=self.enable_storage_metrics,
+                load_cache_event=self.load_cache_event,
+            )
+        else:
+            self.cache_controller = HiCacheController(
+                params.token_to_kv_pool_allocator,
+                self.token_to_kv_pool_host,
+                self.page_size,
+                self.tp_group,
+                load_cache_event=self.load_cache_event,
+                write_policy=server_args.hicache_write_policy,
+                io_backend=server_args.hicache_io_backend,
+                storage_backend=server_args.hicache_storage_backend,
+                prefetch_threshold=prefetch_threshold,
+                model_name=server_args.served_model_name,
+                storage_backend_extra_config=extra_config,
+                pp_rank=self.pp_rank,
+                pp_size=self.pp_size,
+                enable_storage_metrics=self.enable_storage_metrics,
+            )
         self._apply_storage_runtime_config(
             storage_backend=server_args.hicache_storage_backend,
             prefetch_threshold=prefetch_threshold,
@@ -315,12 +332,21 @@ def attach_storage_backend(
             )
 
         try:
-            self.cache_controller.attach_storage_backend(
-                storage_backend=storage_backend,
-                prefetch_threshold=prefetch_threshold,
-                model_name=served_model_name,
-                storage_backend_extra_config=extra_config,
-            )
+            if isinstance(self.cache_controller, HybridCacheController):
+                self.cache_controller.attach_storage_backend(
+                    storage_backend=storage_backend,
+                    prefetch_threshold=prefetch_threshold,
+                    model_name=served_model_name,
+                    storage_backend_extra_config=extra_config,
+                    host_pools=self.cache_controller.mem_pool_host.entries,
+                )
+            else:
+                self.cache_controller.attach_storage_backend(
+                    storage_backend=storage_backend,
+                    prefetch_threshold=prefetch_threshold,
+                    model_name=served_model_name,
+                    storage_backend_extra_config=extra_config,
+                )
         except Exception as e:
             logger.exception(
                 f"Failed to attach storage backend '{storage_backend}': {e}"
@@ -585,6 +611,18 @@ def get_height(self, node: TreeNode):
             height += 1
         return height
 
+    def _get_extra_pools(self) -> dict:
+        if not isinstance(self.cache_controller, HybridCacheController):
+            return {}
+        if isinstance(self.kv_cache, NSATokenToKVPool):
+            pool = PoolTransfer(
+                name=PoolName.INDEXER,
+                hit_policy=PoolHitPolicy.ALL_PAGES,
+            )
+            return {"extra_pools": [pool]}
+        else:
+            return {}
+
     def clear_storage_backend(self) -> bool:
         if self.enable_storage:
             try:
@@ -615,12 +653,14 @@ def write_backup(self, node: TreeNode, write_back=False):
         host_indices = self.cache_controller.write(
             device_indices=node.value,
             node_id=node.id,
+            **self._get_extra_pools(),
         )
         if host_indices is None:
             self.evict_host(len(node.value))
             host_indices = self.cache_controller.write(
                 device_indices=node.value,
                 node_id=node.id,
+                **self._get_extra_pools(),
             )
         if host_indices is not None:
             node.host_value = host_indices.clone()
@@ -642,7 +682,8 @@ def write_backup_storage(self, node: TreeNode):
         )
 
         operation_id = self.cache_controller.write_storage(
-            node.host_value, node.key, node.hash_value, prefix_keys
+            node.host_value, node.key, node.hash_value, prefix_keys,
+            **self._get_extra_pools(),
         )
         self.ongoing_backup[operation_id] = node
         node.protect_host()
@@ -913,12 +954,16 @@ def load_back(
             return None
 
         device_indices = self.cache_controller.load(
-            host_indices=host_indices, node_id=last_hit_node.id
+            host_indices=host_indices,
+            node_id=last_hit_node.id,
+            **self._get_extra_pools(),
         )
         if device_indices is None:
             self.evict(EvictParams(num_tokens=len(host_indices)))
             device_indices = self.cache_controller.load(
-                host_indices=host_indices, node_id=last_hit_node.id
+                host_indices=host_indices,
+                node_id=last_hit_node.id,
+                **self._get_extra_pools(),
             )
         self.dec_lock_ref(ancester_node)
         if device_indices is None:
@@ -1219,7 +1264,12 @@ def prefetch_from_storage(
             # no sufficient host memory for prefetch
             return
         operation = self.cache_controller.prefetch(
-            req_id, host_indices, new_input_tokens, last_hash, prefix_keys
+            req_id,
+            host_indices,
+            new_input_tokens,
+            last_hash,
+            prefix_keys,
+            **self._get_extra_pools(),
         )
         self.ongoing_prefetch[req_id] = (
             last_host_node,

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -163,6 +163,7 @@ def __init__(
         pp_rank: int = 0,
         pp_size: int = 1,
         transfer_layer_num: Optional[int] = None,
+        enable_storage_metrics: bool = False,
     ):
         startup_storage_backend = storage_backend
         super().__init__(
@@ -179,6 +180,7 @@ def __init__(
             storage_backend_extra_config=storage_backend_extra_config,
             pp_rank=pp_rank,
             pp_size=pp_size,
+            enable_storage_metrics=enable_storage_metrics,
         )
         # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)),
         # not just the full attention layers reported by full_kv_pool.
@@ -230,7 +232,10 @@ def write(
         if host_indices is None:
             return None
         pool_transfers = self._resolve_pool_transfers_allocation(
-            extra_pools, alloc_host=True
+            extra_pools,
+            alloc_host=True,
+            kv_device_indices=device_indices,
+            kv_host_indices=host_indices,
         )
         if pool_transfers is None and extra_pools:
             self.mem_pool_host.free(host_indices)
@@ -289,7 +294,10 @@ def load(
             device_indices = torch.empty((0,), dtype=torch.int64, device=self.device)
 
         pool_transfers = self._resolve_pool_transfers_allocation(
-            extra_pools, alloc_host=False
+            extra_pools,
+            alloc_host=False,
+            kv_device_indices=device_indices,
+            kv_host_indices=host_indices,
         )
         if pool_transfers is None and extra_pools:
             if need_load_kv:
@@ -415,6 +423,7 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]:
     def _page_transfer(self, operation):
         # Transfer extra pools
         if operation.pool_transfers and not operation.is_terminated():
+            self._resolve_shared_pool_transfers(operation)
             results = self.storage_backend.batch_get_v2(operation.pool_transfers)
             operation.pool_storage_result.update_extra_pool_hit_pages(results)
 
@@ -424,12 +433,20 @@ def _page_transfer(self, operation):
     def _page_backup(self, operation):
         # Backup extra pools
         if operation.pool_transfers:
+            self._resolve_shared_pool_transfers(operation)
             results = self.storage_backend.batch_set_v2(operation.pool_transfers)
             operation.pool_storage_result.update_extra_pool_hit_pages(results)
 
         # Backup kv pools
         super()._page_backup(operation)
 
+    def _resolve_shared_pool_transfers(self, operation):
+        for transfer in operation.pool_transfers:
+            entry = self.mem_pool_host.entry_map.get(transfer.name)
+            if entry.share_indices_with_anchor:
+                transfer.keys = operation.hash_value
+                transfer.host_indices = operation.host_indices
+
     def _sync_trailing_keys(
         self,
         pool_transfers: list[PoolTransfer],
@@ -454,6 +471,8 @@ def _resolve_pool_transfers_allocation(
         self,
         extra_pools: Optional[list[PoolTransfer]],
         alloc_host: bool,
+        kv_device_indices: Optional[torch.Tensor] = None,
+        kv_host_indices: Optional[torch.Tensor] = None,
     ) -> Optional[list[PoolTransfer]]:
         """Auto-alloc host or device indices for PoolTransfers where they are None."""
         if not extra_pools:
@@ -463,6 +482,10 @@ def _resolve_pool_transfers_allocation(
             entry = self.mem_pool_host.entry_map.get(pool.name)
             if entry is None:
                 continue
+            if entry.share_indices_with_anchor:
+                pool.device_indices = kv_device_indices
+                pool.host_indices = kv_host_indices
+                continue
             if alloc_host:
                 if pool.host_indices is not None or pool.device_indices is None:
                     continue