diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index aece03b44..5a0ff97f9 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 5bedb381c..d6daecf3a 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/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()