From 5fbe43381fdc7539e37070fc65a9b58cfdc9bb11 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 01:00:09 +0800 Subject: [PATCH] feat: add CP HiCache node metadata --- python/sglang/srt/mem_cache/hiradix_cache.py | 48 +++++++++ python/sglang/srt/mem_cache/radix_cache.py | 2 + .../mem_cache/test_cp_hicache_metadata.py | 101 ++++++++++++++++++ 3 files changed, 151 insertions(+) create mode 100644 test/registered/unit/mem_cache/test_cp_hicache_metadata.py diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 1075b0fc1..0ef4017d9 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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): diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 8b5765578..40e8bee43 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py new file mode 100644 index 000000000..798c23021 --- /dev/null +++ b/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()