feat: track CP HiCache write-back state

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -15,6 +15,7 @@ import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
+from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,
     DecLockRefResult,
@@ -27,6 +28,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
     MatchResult,
 )
+from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
     MLATokenToKVPool,
@@ -117,6 +119,21 @@ class HiRadixCache(RadixCache):
             bind_to_closest_numa_node_cuda()
 
         self.page_size = params.page_size
+        self._uses_cp_hicache = isinstance(
+            params.token_to_kv_pool_allocator,
+            CPSharedPagedTokenToKVPoolAllocator,
+        )
+        if self._uses_cp_hicache:
+            if server_args.hicache_storage_backend is not None:
+                raise ValueError(
+                    "CP shared KV HiCache host integration does not support storage backends."
+                )
+            if not isinstance(
+                params.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool
+            ):
+                raise ValueError(
+                    "CP shared KV HiCache host integration requires NSATokenToKVPool."
+                )
         self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
 
         if isinstance(self.kv_cache, MHATokenToKVPool):
@@ -171,6 +188,13 @@ class HiRadixCache(RadixCache):
         self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
 
         self.load_cache_event = threading.Event()
+        cp_shared_kv_layout = None
+        if self._uses_cp_hicache:
+            cp_shared_kv_layout = CpSharedKVLayout(
+                page_size=self.page_size,
+                cp_size=params.token_to_kv_pool_allocator.cp_size,
+                cp_rank=params.token_to_kv_pool_allocator.cp_rank,
+            )
         self.cache_controller = HiCacheController(
             params.token_to_kv_pool_allocator,
             self.token_to_kv_pool_host,
@@ -186,6 +210,7 @@ class HiRadixCache(RadixCache):
             pp_rank=self.pp_rank,
             pp_size=self.pp_size,
             enable_storage_metrics=self.enable_storage_metrics,
+            cp_shared_kv_layout=cp_shared_kv_layout,
         )
         self._apply_storage_runtime_config(
             storage_backend=server_args.hicache_storage_backend,
@@ -679,7 +704,44 @@ class HiRadixCache(RadixCache):
             logger.warning("Hierarchical cache storage backend is not enabled.")
             return False
 
+    def _node_backuped(self, node: TreeNode) -> bool:
+        if self._uses_cp_hicache:
+            return node.host_len > 0 and node.cp_hicache is not None
+        return node.host_value is not None
+
+    def _node_host_len(self, node: TreeNode) -> int:
+        if self._uses_cp_hicache:
+            return node.host_len
+        return len(node.host_value)
+
+    def _node_host_evict_indices(self, node: TreeNode) -> torch.Tensor:
+        if self._uses_cp_hicache:
+            return node.cp_hicache.host_indices
+        return node.host_value
+
     def write_backup(self, node: TreeNode, write_back=False):
+        if self._uses_cp_hicache:
+            result = self.cache_controller.write(
+                device_indices=node.value,
+                node_id=node.id,
+            )
+            if getattr(result, "metadata", None) is None:
+                self._evict_host_for_physical_slots(result.required_host_slots)
+                result = self.cache_controller.write(
+                    device_indices=node.value,
+                    node_id=node.id,
+                )
+            if getattr(result, "metadata", None) is None:
+                return 0
+
+            node.host_len = len(node.value)
+            node.cp_hicache = result.metadata
+            node.host_value = None
+            self.ongoing_write_through[node.id] = node
+            if not write_back:
+                self.inc_node_lock_ref(node)
+            return len(node.cp_hicache.host_indices)
+
         host_indices = self.cache_controller.write(
             device_indices=node.value,
             node_id=node.id,
@@ -725,7 +787,7 @@ class HiRadixCache(RadixCache):
             return
         node.hit_count += 1
 
-        if not node.backuped:
+        if not self._node_backuped(node):
             if node.hit_count >= self.write_through_threshold:
                 # write to host if the node is not backuped
                 self.write_backup(node)