From 068aa2f91000132f92193bf93260866a64d9d856 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 01:29:39 +0800 Subject: [PATCH] feat: track CP HiCache write-back state --- python/sglang/srt/mem_cache/hiradix_cache.py | 64 ++++++++++++- .../mem_cache/test_cp_hicache_metadata.py | 90 ++++++++++++++++++- 2 files changed, 152 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b57ec8775..94a722384 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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) 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 a237c9b9e..13e06cd89 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,8 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"): if _mod not in sys.modules: sys.modules[_mod] = MagicMock() -from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata +from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache +from sglang.srt.mem_cache.radix_cache import TreeNode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -139,5 +140,92 @@ class TestCpHiCacheNodeMetadata(CustomTestCase): ) +class FakeWriteFailure: + metadata = None + + def __init__(self, required_host_slots): + self.required_host_slots = required_host_slots + + +class FakeWriteSuccess: + required_host_slots = 0 + + def __init__(self, metadata): + self.metadata = metadata + + +class FakeWriteController: + def __init__(self, required_host_slots): + self.required_host_slots = required_host_slots + self.calls = 0 + self.write_policy = "write_through" + + def write(self, device_indices, node_id=-1, priority=None): + self.calls += 1 + if self.calls == 1: + return FakeWriteFailure(self.required_host_slots) + return FakeWriteSuccess( + CpHiCacheNodeMetadata( + logical_len=len(device_indices), + owned_positions=torch.tensor([0], dtype=torch.int64), + host_indices=torch.tensor([99], dtype=torch.int64), + ) + ) + + +class TestHiRadixCacheCPBackup(CustomTestCase): + def test_node_backuped_uses_cp_metadata(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + node = TreeNode() + node.host_len = 8 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=8, + owned_positions=torch.tensor([1, 2], dtype=torch.int64), + host_indices=torch.tensor([10, 11], dtype=torch.int64), + ) + + self.assertTrue(cache._node_backuped(node)) + + def test_inc_hit_count_does_not_rewrite_cp_backed_node(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.write_through_threshold = 1 + cache.cache_controller = type( + "Controller", (), {"write_policy": "write_through"} + )() + cache.write_backup = lambda node: (_ for _ in ()).throw( + AssertionError("must not rewrite") + ) + node = TreeNode() + node.host_len = 4 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.tensor([], dtype=torch.int64), + host_indices=torch.tensor([], dtype=torch.int64), + ) + + cache._inc_hit_count(node) + + self.assertEqual(node.hit_count, 1) + + def test_write_backup_retries_by_required_physical_slots(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = FakeWriteController(required_host_slots=3) + cache._evict_host_for_physical_slots = lambda required: setattr( + cache, "evicted_required", required + ) + cache.ongoing_write_through = {} + cache.inc_node_lock_ref = lambda node: None + node = TreeNode() + node.value = torch.arange(16, dtype=torch.int64) + + cache.write_backup(node) + + self.assertEqual(cache.evicted_required, 3) + self.assertEqual(node.host_len, 16) + + if __name__ == "__main__": unittest.main()