feat: track CP HiCache write-back state

2026-05-08 01:29:39 +08:00
parent a8816f18dd
commit 068aa2f910
2 changed files with 152 additions and 2 deletions
@@ -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()