feat: restore CP HiCache host hits as logical locs

2026-05-08 02:34:16 +08:00
parent f10a8983c3
commit 829f3ebceb
2 changed files with 195 additions and 6 deletions
@@ -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()