feat: restore CP HiCache host hits as logical locs

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -1180,6 +1180,71 @@ class HiRadixCache(RadixCache):
         self, node: TreeNode, mem_quota: Optional[int] = None
     ) -> Optional[torch.Tensor]:
 
+        if self._uses_cp_hicache:
+            start_time = time.perf_counter()
+            last_hit_node = node
+            nodes_to_load = []
+            while node.evicted:
+                assert self._node_backuped(
+                    node
+                ), "No backup available on evicted nodes, should not happen"
+                nodes_to_load.insert(0, node)
+                node = node.parent
+            else:
+                ancester_node = node
+
+            # protect the ancestor nodes from eviction
+            result = self.inc_lock_ref(ancester_node)
+            delta = result.delta
+
+            # load it all or not at all
+            host_hit_len = sum(self._node_host_len(n) for n in nodes_to_load)
+            if host_hit_len < self.load_back_threshold or (
+                host_hit_len > mem_quota + delta if mem_quota is not None else False
+            ):
+                # skip loading back if the total size is too small or exceeding the memory quota
+                self.dec_lock_ref(ancester_node)
+                return None
+
+            device_indices = self.cache_controller.load_cp(
+                nodes_to_load, node_id=last_hit_node.id
+            )
+            if device_indices is None:
+                self.evict(EvictParams(num_tokens=host_hit_len))
+                device_indices = self.cache_controller.load_cp(
+                    nodes_to_load, node_id=last_hit_node.id
+                )
+            self.dec_lock_ref(ancester_node)
+            if device_indices is None:
+                # no sufficient GPU memory to load back KV caches
+                logger.warning(
+                    "load_back: FAILED to load %d tokens for node %d "
+                    "even after eviction (evictable_size=%d)",
+                    host_hit_len,
+                    last_hit_node.id,
+                    self.evictable_size_,
+                )
+                return None
+
+            self.ongoing_load_back[last_hit_node.id] = last_hit_node
+            offset = 0
+            for loaded_node in nodes_to_load:
+                host_len = self._node_host_len(loaded_node)
+                loaded_node.value = device_indices[offset : offset + host_len].clone()
+                offset += host_len
+            self.evictable_size_ += len(device_indices)
+            self.inc_lock_ref(last_hit_node)
+
+            if self.metrics_collector is not None:
+                self.metrics_collector.observe_load_back_duration(
+                    time.perf_counter() - start_time
+                )
+                self.metrics_collector.increment_load_back_num_tokens(
+                    len(device_indices)
+                )
+
+            return device_indices
+
         start_time = time.perf_counter()
         last_hit_node = node
         nodes_to_load = []
@@ -1464,11 +1529,23 @@ class HiRadixCache(RadixCache):
 
         host_hit_length = 0
         last_host_node = last_node
-        while last_node.evicted:
-            host_hit_length += len(last_node.host_value)
-            last_node = last_node.parent
-        while not last_host_node.backuped:
-            last_host_node = last_host_node.parent
+        if self._uses_cp_hicache:
+            while last_node.evicted:
+                host_hit_length += self._node_host_len(last_node)
+                last_node = last_node.parent
+            while (
+                last_host_node != self.root_node
+                and not self._node_backuped(last_host_node)
+            ):
+                last_host_node = last_host_node.parent
+            if not self._node_backuped(last_host_node):
+                last_host_node = self.root_node
+        else:
+            while last_node.evicted:
+                host_hit_length += len(last_node.host_value)
+                last_node = last_node.parent
+            while not last_host_node.backuped:
+                last_host_node = last_host_node.parent
 
         return MatchResult(
             device_indices=value,

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -10,7 +10,7 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
     if _mod not in sys.modules:
         sys.modules[_mod] = MagicMock()
 
-from sglang.srt.mem_cache.base_prefix_cache import EvictParams
+from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
 from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
 from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -498,5 +498,117 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         self.assertEqual(parent.cp_hicache.host_indices.tolist(), [80])
 
 
+class TestHiRadixCacheCPLoadBack(CustomTestCase):
+    def test_cp_load_back_uses_host_len_not_host_value(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.root_node = TreeNode()
+        cache.device = "cpu"
+        cache.load_back_threshold = 1
+        cache.evictable_size_ = 0
+        cache.metrics_collector = None
+        cache.ongoing_load_back = {}
+        cache.cache_controller = type(
+            "Controller",
+            (),
+            {
+                "load_cp": lambda self, nodes, node_id=-1: torch.arange(
+                    32, 40, dtype=torch.int64
+                )
+            },
+        )()
+        cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
+        cache.dec_lock_ref = lambda node: None
+        cache.evict = lambda params: None
+        parent = cache.root_node
+        parent.key = RadixKey([])
+        parent.value = torch.empty((0,), dtype=torch.int64)
+        node = TreeNode()
+        node.parent = parent
+        node.key = RadixKey(list(range(8)))
+        node.value = None
+        node.host_value = None
+        node.host_len = 8
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=8,
+            owned_positions=torch.tensor([0, 1], dtype=torch.int64),
+            host_indices=torch.tensor([50, 51], dtype=torch.int64),
+        )
+
+        loaded = cache.load_back(node)
+
+        self.assertEqual(loaded.tolist(), list(range(32, 40)))
+        self.assertEqual(node.value.tolist(), list(range(32, 40)))
+
+    def test_cp_load_back_threshold_uses_logical_length(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.root_node = TreeNode()
+        cache.root_node.value = torch.empty((0,), dtype=torch.int64)
+        cache.load_back_threshold = 5
+        cache.evictable_size_ = 0
+        cache.metrics_collector = None
+        cache.ongoing_load_back = {}
+        cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
+        cache.dec_lock_ref = lambda node: None
+        cache.cache_controller = type(
+            "Controller",
+            (),
+            {
+                "load_cp": lambda self, nodes, node_id=-1: torch.arange(
+                    10, 16, dtype=torch.int64
+                )
+            },
+        )()
+        node = TreeNode()
+        node.parent = cache.root_node
+        node.value = None
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            owned_positions=torch.empty((0,), dtype=torch.int64),
+            host_indices=torch.empty((0,), dtype=torch.int64),
+        )
+
+        loaded = cache.load_back(node)
+
+        self.assertEqual(loaded.tolist(), [10, 11, 12, 13, 14, 15])
+
+    def test_cp_match_prefix_counts_logical_host_hit_without_host_value(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.device = "cpu"
+        cache.disable = False
+        cache.page_size = 1
+        cache.get_child_key_fn = lambda key: key.token_ids[0]
+        cache.key_match_fn = lambda child_key, key: sum(
+            1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
+        )
+        cache.maybe_bigram_convert = lambda key: (key, None)
+        root = TreeNode()
+        root.key = RadixKey([])
+        root.value = torch.empty((0,), dtype=torch.int64)
+        root.host_len = 0
+        cache.root_node = root
+        node = TreeNode()
+        node.parent = root
+        node.key = RadixKey(list(range(8)))
+        node.value = None
+        node.host_value = None
+        node.host_len = 8
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=8,
+            owned_positions=torch.tensor([0, 1], dtype=torch.int64),
+            host_indices=torch.tensor([50, 51], dtype=torch.int64),
+        )
+        root.children[0] = node
+
+        result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8)))))
+
+        self.assertEqual(result.host_hit_length, 8)
+        self.assertIs(result.last_host_node, node)
+        self.assertIs(result.last_device_node, root)
+
+
 if __name__ == "__main__":
     unittest.main()