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