feat: add CP HiCache node metadata

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -7,6 +7,7 @@ import logging
 import os
 import threading
 import time
+from dataclasses import dataclass
 from queue import Empty
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
@@ -54,6 +55,53 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class CpHiCacheNodeMetadata:
+    logical_len: int
+    owned_positions: torch.Tensor
+    host_indices: torch.Tensor
+
+    def __post_init__(self):
+        if self.logical_len < 0:
+            raise ValueError(f"logical_len must be non-negative, got {self.logical_len}")
+        self.owned_positions = self.owned_positions.to(device="cpu", dtype=torch.int64)
+        self.host_indices = self.host_indices.to(device="cpu", dtype=torch.int64)
+        if self.owned_positions.numel() != self.host_indices.numel():
+            raise ValueError(
+                "owned_positions and host_indices must have same length, got "
+                f"{self.owned_positions.numel()} and {self.host_indices.numel()}"
+            )
+        if self.owned_positions.numel() > 0:
+            if torch.any(self.owned_positions < 0) or torch.any(
+                self.owned_positions >= self.logical_len
+            ):
+                raise ValueError("owned_positions must be in [0, logical_len)")
+            if torch.any(self.owned_positions[1:] < self.owned_positions[:-1]):
+                raise ValueError("owned_positions must be sorted")
+
+    def split(
+        self, split_len: int
+    ) -> tuple["CpHiCacheNodeMetadata", "CpHiCacheNodeMetadata"]:
+        if split_len < 0 or split_len > self.logical_len:
+            raise ValueError(
+                f"split_len must be in [0, {self.logical_len}], got {split_len}"
+            )
+        parent_mask = self.owned_positions < split_len
+        child_mask = ~parent_mask
+        return (
+            CpHiCacheNodeMetadata(
+                logical_len=split_len,
+                owned_positions=self.owned_positions[parent_mask],
+                host_indices=self.host_indices[parent_mask],
+            ),
+            CpHiCacheNodeMetadata(
+                logical_len=self.logical_len - split_len,
+                owned_positions=self.owned_positions[child_mask] - split_len,
+                host_indices=self.host_indices[child_mask],
+            ),
+        )
+
+
 class HiRadixCache(RadixCache):
 
     def __init__(self, params: CacheInitParams, server_args: ServerArgs):

python/sglang/srt/mem_cache/radix_cache.py

@@ -140,6 +140,8 @@ class TreeNode:
         self.host_ref_counter = 0
         # store the host indices of KV cache
         self.host_value: Optional[torch.Tensor] = None
+        self.host_len: int = 0
+        self.cp_hicache = None
         # store hash values of each pages
         self.hash_value: Optional[List[str]] = None
         # priority for priority-aware eviction

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -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()