feat: add CP HiCache node metadata

2026-05-08 01:00:09 +08:00
parent da4994ec26
commit 5fbe43381f
3 changed files with 151 additions and 0 deletions
@@ -0,0 +1,101 @@
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+# Stub out sgl_kernel before any sglang import so this CPU unit test does not
+# require CUDA extension libraries to be installed.
+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, HiRadixCache
+from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
+
+
+class TestCpHiCacheNodeMetadata(CustomTestCase):
+    def test_split_zero_len_moves_all_positions_to_child(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=8,
+            owned_positions=torch.tensor([1, 3, 7], dtype=torch.int64),
+            host_indices=torch.tensor([10, 11, 12], dtype=torch.int64),
+        )
+
+        parent, child = metadata.split(0)
+
+        self.assertEqual(parent.logical_len, 0)
+        self.assertEqual(parent.owned_positions.tolist(), [])
+        self.assertEqual(parent.host_indices.tolist(), [])
+        self.assertEqual(child.logical_len, 8)
+        self.assertEqual(child.owned_positions.tolist(), [1, 3, 7])
+        self.assertEqual(child.host_indices.tolist(), [10, 11, 12])
+
+    def test_split_rebases_child_positions(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=10,
+            owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
+            host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
+        )
+
+        parent, child = metadata.split(5)
+
+        self.assertEqual(parent.logical_len, 5)
+        self.assertEqual(parent.owned_positions.tolist(), [0, 2])
+        self.assertEqual(parent.host_indices.tolist(), [20, 21])
+        self.assertEqual(child.logical_len, 5)
+        self.assertEqual(child.owned_positions.tolist(), [0, 4])
+        self.assertEqual(child.host_indices.tolist(), [22, 23])
+
+    def test_zero_owned_metadata_is_valid(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=64,
+            owned_positions=torch.empty((0,), dtype=torch.int64),
+            host_indices=torch.empty((0,), dtype=torch.int64),
+        )
+
+        self.assertEqual(metadata.logical_len, 64)
+        self.assertEqual(metadata.owned_positions.device.type, "cpu")
+        self.assertEqual(metadata.host_indices.device.type, "cpu")
+
+    def test_invalid_split_raises(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=4,
+            owned_positions=torch.tensor([1], dtype=torch.int64),
+            host_indices=torch.tensor([9], dtype=torch.int64),
+        )
+
+        with self.assertRaisesRegex(ValueError, "split_len"):
+            metadata.split(5)
+
+    def test_unsorted_positions_raise(self):
+        with self.assertRaisesRegex(ValueError, "sorted"):
+            CpHiCacheNodeMetadata(
+                logical_len=4,
+                owned_positions=torch.tensor([2, 1], dtype=torch.int64),
+                host_indices=torch.tensor([9, 10], dtype=torch.int64),
+            )
+
+    def test_length_mismatch_raises(self):
+        with self.assertRaisesRegex(ValueError, "same length"):
+            CpHiCacheNodeMetadata(
+                logical_len=4,
+                owned_positions=torch.tensor([1, 2], dtype=torch.int64),
+                host_indices=torch.tensor([9], dtype=torch.int64),
+            )
+
+    def test_out_of_range_positions_raise(self):
+        with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
+            CpHiCacheNodeMetadata(
+                logical_len=4,
+                owned_positions=torch.tensor([4], dtype=torch.int64),
+                host_indices=torch.tensor([9], dtype=torch.int64),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()