import sys import types import unittest from unittest.mock import patch import torch # Prefer the real sgl_kernel package when the test image provides it so custom # Torch operators are registered. Fall back to stubs on local CPU-only hosts. try: import sgl_kernel # noqa: F401 import sgl_kernel.kvcacheio # noqa: F401 except (ImportError, RuntimeError): if "sgl_kernel" not in sys.modules: sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel") sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py" sys.modules["sgl_kernel"].__path__ = [] if not hasattr(sys.modules["sgl_kernel"], "__getattr__"): def _sgl_kernel_getattr(name): if name.startswith("__"): raise AttributeError(name) fn = lambda *args, **kwargs: None setattr(sys.modules["sgl_kernel"], name, fn) return fn sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr if "sgl_kernel.quantization" not in sys.modules: quantization_stub = types.ModuleType("sgl_kernel.quantization") quantization_stub.__file__ = "sgl_kernel_quantization_stub.py" def _quantization_getattr(name): if name.startswith("__"): raise AttributeError(name) fn = lambda *args, **kwargs: None setattr(quantization_stub, name, fn) return fn quantization_stub.__getattr__ = _quantization_getattr for _name in ( "ggml_dequantize", "ggml_moe_a8", "ggml_moe_a8_vec", "ggml_moe_get_block_size", "ggml_mul_mat_a8", "ggml_mul_mat_vec_a8", ): setattr(quantization_stub, _name, lambda *args, **kwargs: None) sys.modules["sgl_kernel.quantization"] = quantization_stub for _name in ( "sgl_per_token_group_quant_8bit", "sgl_per_token_group_quant_fp8", "sgl_per_token_quant_fp8", "fp8_blockwise_scaled_mm", "fp8_scaled_mm", "silu_and_mul", ): if not hasattr(sys.modules["sgl_kernel"], _name): setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None) _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") for _schema in ( "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", ): try: _sgl_kernel_lib.define(_schema) except RuntimeError as exc: if ( "already" not in str(exc).lower() and "duplicate" not in str(exc).lower() ): raise if "sgl_kernel.kvcacheio" not in sys.modules: sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio") _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") for _schema in ( "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", ): try: _sgl_kernel_lib.define(_schema) except RuntimeError as exc: if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): raise from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams from sglang.srt.mem_cache.hiradix_cache import ( CpHiCacheNodeMetadata, HiRadixCache, _compute_shared_hicache_token_capacities, ) 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 TestCpHiCacheImports(CustomTestCase): def test_cp_hicache_public_imports_without_sgl_kernel(self): import subprocess subprocess.run( [ sys.executable, "-c", "from sglang.srt.mem_cache.hiradix_cache import HiRadixCache, CpHiCacheNodeMetadata; " "from sglang.srt.managers.cache_controller import HiCacheController; " "print('OK')", ], check=True, capture_output=True, text=True, ) class TestHiRadixCacheCPDraftHostPool(CustomTestCase): def test_shared_budget_keeps_draft_at_least_target_capacity(self): target_tokens, draft_tokens = _compute_shared_hicache_token_capacities( total_host_bytes=1000, target_size_per_token=6, draft_size_per_token=2, page_size=10, ) self.assertEqual(target_tokens, 120) self.assertEqual(draft_tokens, 140) self.assertGreaterEqual(draft_tokens, target_tokens) self.assertLessEqual(target_tokens * 6 + draft_tokens * 2, 1000) def test_shared_budget_handles_equal_target_and_draft_size(self): target_tokens, draft_tokens = _compute_shared_hicache_token_capacities( total_host_bytes=1000, target_size_per_token=6, draft_size_per_token=6, page_size=10, ) self.assertEqual(target_tokens, 80) self.assertEqual(draft_tokens, 80) self.assertLessEqual(target_tokens * 6 + draft_tokens * 6, 1000) def test_reset_clears_target_and_draft_host_pools(self): class ClearablePool: def __init__(self): self.clear_calls = 0 def clear(self): self.clear_calls += 1 class Controller: def __init__(self): self.reset_calls = 0 self.clear_draft_calls = 0 def reset(self): self.reset_calls += 1 def clear_draft_host_pool(self): self.clear_draft_calls += 1 target_pool = ClearablePool() controller = Controller() cache = HiRadixCache.__new__(HiRadixCache) cache.cache_controller = controller cache.token_to_kv_pool_host = target_pool cache.prefetch_loaded_tokens_by_reqid = {} cache.evictable_host_leaves = set() cache.pinned_size_ = 1 cache.evictable_leaves = set() cache._record_all_cleared_event = lambda: None cache.reset() self.assertEqual(controller.reset_calls, 1) self.assertEqual(target_pool.clear_calls, 1) self.assertEqual(controller.clear_draft_calls, 1) 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), page_owners=torch.zeros(max(8, 0), dtype=torch.int8), page_size=1, ) 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), page_owners=torch.zeros(max(10, 0), dtype=torch.int8), page_size=1, ) 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_split_keeps_draft_host_indices_aligned_with_owned_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), draft_host_indices=torch.tensor([120, 121, 122, 123], dtype=torch.int64), page_owners=torch.zeros(max(10, 0), dtype=torch.int8), page_size=1, ) parent, child = metadata.split(5) self.assertEqual(parent.host_indices.tolist(), [20, 21]) self.assertEqual(parent.draft_host_indices.tolist(), [120, 121]) self.assertEqual(child.owned_positions.tolist(), [0, 4]) self.assertEqual(child.host_indices.tolist(), [22, 23]) self.assertEqual(child.draft_host_indices.tolist(), [122, 123]) def test_zero_owned_metadata_is_valid(self): metadata = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.empty((0,), dtype=torch.int32), host_indices=torch.empty((0,), dtype=torch.int32), page_owners=torch.zeros(max(64, 0), dtype=torch.int8), page_size=1, ) self.assertEqual(metadata.logical_len, 64) self.assertEqual(metadata.owned_positions.device.type, "cpu") self.assertEqual(metadata.host_indices.device.type, "cpu") self.assertEqual(metadata.owned_positions.dtype, torch.int64) self.assertEqual(metadata.host_indices.dtype, torch.int64) def test_non_int64_inputs_are_converted(self): metadata = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([1, 3], dtype=torch.int32), host_indices=torch.tensor([10, 11], dtype=torch.int32), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) self.assertEqual(metadata.owned_positions.dtype, torch.int64) self.assertEqual(metadata.host_indices.dtype, torch.int64) def test_negative_logical_len_raises(self): with self.assertRaisesRegex(ValueError, "logical_len"): CpHiCacheNodeMetadata( logical_len=-1, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(0, dtype=torch.int8), page_size=1, ) def test_metadata_does_not_alias_input_tensors(self): owned_positions = torch.tensor([1, 3], dtype=torch.int64) host_indices = torch.tensor([10, 11], dtype=torch.int64) metadata = CpHiCacheNodeMetadata( logical_len=4, owned_positions=owned_positions, host_indices=host_indices, page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) owned_positions[0] = 2 host_indices[0] = 12 self.assertEqual(metadata.owned_positions.tolist(), [1, 3]) self.assertEqual(metadata.host_indices.tolist(), [10, 11]) 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), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) 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), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) def test_duplicate_positions_raise(self): with self.assertRaisesRegex(ValueError, "strictly increasing"): CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([1, 1], dtype=torch.int64), host_indices=torch.tensor([9, 10], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) 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), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) def test_draft_host_length_mismatch_raises(self): with self.assertRaisesRegex(ValueError, "draft_host_indices.*same length"): CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([1, 2], dtype=torch.int64), host_indices=torch.tensor([9, 10], dtype=torch.int64), draft_host_indices=torch.tensor([109], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) 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), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) # ── New: validators for page_owners + page_size (the fields that carry the # CP owner pattern across a HiCache write→load round-trip). def test_zero_page_size_raises(self): with self.assertRaisesRegex(ValueError, "page_size must be positive"): CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.empty((0,), dtype=torch.int8), page_size=0, ) def test_logical_len_not_multiple_of_page_size_raises(self): with self.assertRaisesRegex(ValueError, "multiple of"): CpHiCacheNodeMetadata( logical_len=10, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(2, dtype=torch.int8), page_size=4, # 10 % 4 != 0 ) def test_page_owners_length_mismatch_raises(self): with self.assertRaisesRegex(ValueError, "page_owners length"): CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(3, dtype=torch.int8), # expected 8/4=2 page_size=4, ) def test_negative_page_owners_raises(self): with self.assertRaisesRegex(ValueError, "page_owners.*non-negative"): CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.tensor([-1, 0], dtype=torch.int8), page_size=4, ) def test_page_owners_normalized_to_int8_cpu(self): metadata = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), # Pass int64 to test normalization. page_owners=torch.tensor([0, 1], dtype=torch.int64), page_size=4, ) self.assertEqual(metadata.page_owners.dtype, torch.int8) self.assertEqual(metadata.page_owners.device.type, "cpu") self.assertEqual(metadata.page_owners.tolist(), [0, 1]) def test_page_owners_not_aliased_to_input(self): page_owners = torch.tensor([0, 1, 0, 1], dtype=torch.int8) metadata = CpHiCacheNodeMetadata( logical_len=16, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=page_owners, page_size=4, ) page_owners[0] = 1 self.assertEqual(metadata.page_owners.tolist(), [0, 1, 0, 1]) def test_split_non_page_aligned_split_len_raises(self): metadata = CpHiCacheNodeMetadata( logical_len=16, owned_positions=torch.tensor([0, 8], dtype=torch.int64), host_indices=torch.tensor([20, 21], dtype=torch.int64), page_owners=torch.zeros(4, dtype=torch.int8), page_size=4, ) with self.assertRaisesRegex(ValueError, "must be a multiple of page_size"): metadata.split(5) # 5 % 4 != 0 def test_split_preserves_page_owners_slice(self): metadata = CpHiCacheNodeMetadata( logical_len=32, # 8 pages of page_size=4 owned_positions=torch.tensor([0, 4, 16, 28], dtype=torch.int64), host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64), page_owners=torch.tensor( [0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.int8 ), page_size=4, ) parent, child = metadata.split(16) # split at page 4 of 8. self.assertEqual(parent.page_owners.tolist(), [0, 1, 0, 1]) self.assertEqual(child.page_owners.tolist(), [0, 1, 0, 1]) self.assertEqual(parent.page_size, 4) self.assertEqual(child.page_size, 4) def test_split_at_zero_yields_empty_parent_page_owners(self): metadata = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.tensor([0], dtype=torch.int64), host_indices=torch.tensor([50], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) parent, child = metadata.split(0) self.assertEqual(parent.page_owners.numel(), 0) self.assertEqual(child.page_owners.tolist(), [0, 1]) def test_split_at_logical_len_yields_empty_child_page_owners(self): metadata = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.tensor([0], dtype=torch.int64), host_indices=torch.tensor([50], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) parent, child = metadata.split(8) self.assertEqual(parent.page_owners.tolist(), [0, 1]) self.assertEqual(child.page_owners.numel(), 0) 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" self.evicted_host_indices = [] 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), page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8), page_size=1, ) ) def evict_host(self, host_indices): self.evicted_host_indices.append(host_indices.clone()) return len(host_indices) class FakeZeroOwnedWriteController: write_policy = "write_through" def write(self, device_indices, node_id=-1, priority=None): return FakeWriteSuccess( CpHiCacheNodeMetadata( logical_len=len(device_indices), owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8), page_size=1, ) ) class FakeEvictionStrategy: def get_priority(self, node): return 0 class FakeEvictDeviceController: write_policy = "write_through" def evict_device(self, device_indices): return len(device_indices) class FakeTokenAllocator: def available_size(self): return 0 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), page_owners=torch.zeros(max(8, 0), dtype=torch.int8), page_size=1, ) self.assertTrue(cache._node_backuped(node)) def test_node_backuped_rejects_missing_draft_metadata_when_draft_attached(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace( has_draft_hicache=True, cp_shared_kv_layout=types.SimpleNamespace(cp_rank=2), ) node = TreeNode() node.id = 123 node.host_len = 64 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) with self.assertRaisesRegex( RuntimeError, "node_id=123.*cp_rank=2.*missing draft_host_indices" ): cache._node_backuped(node) def test_node_backuped_rejects_cp_host_len_without_metadata(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=types.SimpleNamespace(cp_rank=1) ) node = TreeNode() node.id = 124 node.host_len = 64 node.cp_hicache = None with self.assertRaisesRegex( RuntimeError, "node_id=124.*host_len=64.*cp_rank=1.*without cp_hicache" ): cache._node_backuped(node) def test_node_backuped_accepts_empty_draft_metadata_for_zero_owned_rank(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True) node = TreeNode() node.host_len = 64 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), draft_host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) self.assertTrue(cache._node_backuped(node)) def test_node_backuped_excludes_inflight_cp_write_until_ack(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True) cache.ongoing_write_through = {} node = TreeNode() node.id = 125 node.host_len = 64 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), draft_host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) cache.ongoing_write_through[node.id] = node self.assertFalse(cache._node_backuped(node)) cache.ongoing_write_through.clear() self.assertTrue(cache._node_backuped(node)) def test_single_node_write_lock_updates_device_evictable_leaf_set(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_leaves = set() cache.evictable_size_ = 4 cache.protected_size_ = 0 node = TreeNode() node.parent = cache.root_node node.key = RadixKey([1, 2, 3, 4]) node.value = torch.arange(4, dtype=torch.int64) cache.root_node.children[1] = node cache.evictable_leaves.add(node) cache.inc_node_lock_ref(node) self.assertNotIn(node, cache.evictable_leaves) self.assertEqual(cache.evictable_size(), 0) self.assertEqual(cache.protected_size(), 4) cache.dec_node_lock_ref(node) self.assertIn(node, cache.evictable_leaves) self.assertEqual(cache.evictable_size(), 4) self.assertEqual(cache.protected_size(), 0) 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), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache._inc_hit_count(node) self.assertEqual(node.hit_count, 1) def test_inc_hit_count_does_not_duplicate_inflight_cp_write(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.write_through_threshold = 1 cache.ongoing_write_through = {} cache.cache_controller = type( "Controller", (), {"write_policy": "write_through"} )() cache.write_backup = lambda node: (_ for _ in ()).throw( AssertionError("must not launch a second write") ) node = TreeNode() node.id = 127 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), draft_host_indices=torch.tensor([], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache.ongoing_write_through[node.id] = node 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=1) cache.evictable_host_leaves = set() cache.eviction_strategy = FakeEvictionStrategy() cache.get_child_key_fn = lambda key: key.token_ids[0] cache._record_remove_event = lambda node: None cache.ongoing_write_through = {} cache.inc_node_lock_ref = lambda node: None root = TreeNode() root.key = RadixKey(token_ids=[], extra_key=None) root.value = [] cache.root_node = root evictable_node = TreeNode() evictable_node.parent = root evictable_node.key = RadixKey(token_ids=[1], extra_key=None) evictable_node.value = None evictable_node.host_len = 4 evictable_node.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([0], dtype=torch.int64), host_indices=torch.tensor([55], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) root.children[1] = evictable_node cache.evictable_host_leaves.add(evictable_node) node = TreeNode() node.value = torch.arange(16, dtype=torch.int64) cache.write_backup(node) self.assertEqual( cache.cache_controller.evicted_host_indices[0].tolist(), [55] ) self.assertEqual(evictable_node.host_len, 0) self.assertIsNone(evictable_node.cp_hicache) self.assertNotIn(1, root.children) self.assertEqual(node.host_len, 16) def test_write_backup_cp_success_returns_logical_length_for_zero_owned_rank(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = FakeZeroOwnedWriteController() cache.ongoing_write_through = {} cache.inc_node_lock_ref = lambda node: None node = TreeNode() node.id = 123 node.value = torch.arange(16, dtype=torch.int64) backed_len = cache.write_backup(node) self.assertEqual(backed_len, 16) self.assertEqual(node.host_len, 16) self.assertEqual(node.cp_hicache.host_indices.tolist(), []) def test_attach_storage_backend_rejects_cp_hicache_without_controller_call(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = type( "Controller", (), { "attach_storage_backend": lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("controller attach must not be called") ) }, )() ok, message = cache.attach_storage_backend("mooncake") self.assertFalse(ok) self.assertIn("CP shared KV", message) self.assertIn("storage backend", message) def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = FakeEvictDeviceController() cache.evictable_leaves = set() cache.evictable_host_leaves = set() cache.eviction_strategy = FakeEvictionStrategy() cache.token_to_kv_pool_allocator = FakeTokenAllocator() cache.evictable_size_ = 4 cache.protected_size_ = 0 cache._record_remove_event = lambda node: (_ for _ in ()).throw( AssertionError("must not delete backed radix child") ) root = TreeNode() root.key = RadixKey(token_ids=[], extra_key=None) root.value = [] cache.root_node = root node = TreeNode() node.parent = root node.key = RadixKey(token_ids=[1, 2, 3, 4], extra_key=None) node.value = torch.arange(4, dtype=torch.int64) node.host_len = 4 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([0], dtype=torch.int64), host_indices=torch.tensor([55], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) root.children[1] = node cache.evictable_leaves.add(node) cache.evict(EvictParams(num_tokens=4)) self.assertIn(1, root.children) self.assertIsNone(node.value) self.assertIsNotNone(node.cp_hicache) self.assertEqual(node.cp_hicache.host_indices.tolist(), [55]) class TestHiRadixCacheCPSplitEvict(CustomTestCase): def test_split_node_splits_cp_metadata_by_owned_positions(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.get_child_key_fn = lambda key: key.token_ids[0] cache.page_size = 1 root = TreeNode() root.key = RadixKey([]) child = TreeNode() child.parent = root child.key = RadixKey(list(range(10))) child.value = None child.hash_value = None child.host_len = 10 child.cp_hicache = 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), page_owners=torch.zeros(max(10, 0), dtype=torch.int8), page_size=1, ) root.children[0] = child new_node = cache._split_node(child.key, child, 5) self.assertEqual(new_node.host_len, 5) self.assertEqual(child.host_len, 5) self.assertEqual(new_node.cp_hicache.owned_positions.tolist(), [0, 2]) self.assertEqual(new_node.cp_hicache.host_indices.tolist(), [20, 21]) self.assertEqual(child.cp_hicache.owned_positions.tolist(), [0, 4]) self.assertEqual(child.cp_hicache.host_indices.tolist(), [22, 23]) def test_cp_host_eviction_uses_physical_freed_slots_for_progress(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() cache.get_child_key_fn = lambda key: key.token_ids[0] cache.eviction_strategy = type( "Strategy", (), {"get_priority": lambda self, node: 0} )() cache._clear_pin = lambda node: None cache._record_remove_event = lambda node: None cache._update_host_leaf_status = lambda node: None cache._node_host_evict_indices = lambda node: torch.tensor( [99], dtype=torch.int64 ) freed = [] cache.cache_controller = type( "Controller", (), { "evict_host": lambda self, indices: freed.append(indices.clone()) or len(indices) }, )() node = TreeNode() node.parent = cache.root_node node.key = RadixKey([1, 2, 3, 4]) node.value = None node.host_len = 4 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([1], dtype=torch.int64), host_indices=torch.tensor([70], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache.root_node.children[1] = node cache.evictable_host_leaves.add(node) physical_freed = cache._evict_host_for_physical_slots(1) self.assertEqual(physical_freed, 1) self.assertEqual(freed[0].tolist(), [99]) self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) def test_cp_host_eviction_unlinks_stale_leaf_to_free_parent_slots(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() cache.get_child_key_fn = lambda key: key.token_ids[0] cache.eviction_strategy = type( "Strategy", (), {"get_priority": lambda self, node: 0} )() cache._clear_pin = lambda node: None cache._record_remove_event = lambda node: None freed = [] cache.cache_controller = type( "Controller", (), { "evict_host": lambda self, indices: freed.append(indices.clone()) or len(indices) }, )() parent = TreeNode() parent.parent = cache.root_node parent.key = RadixKey([1]) parent.value = None parent.host_len = 4 parent.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([2], dtype=torch.int64), host_indices=torch.tensor([80], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache.root_node.children[1] = parent stale_child = TreeNode() stale_child.parent = parent stale_child.key = RadixKey([2]) stale_child.value = None stale_child.host_len = 0 stale_child.cp_hicache = None parent.children[2] = stale_child cache.evictable_host_leaves.add(stale_child) physical_freed = cache._evict_host_for_physical_slots(1) self.assertEqual(physical_freed, 1) self.assertEqual(freed[0].tolist(), [80]) self.assertNotIn(2, parent.children) self.assertEqual(parent.host_len, 0) self.assertIsNone(parent.cp_hicache) self.assertNotIn(1, cache.root_node.children) def test_cp_host_eviction_preserves_parent_with_sibling_after_stale_cleanup( self, ): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() cache.get_child_key_fn = lambda key: key.token_ids[0] cache.eviction_strategy = type( "Strategy", (), {"get_priority": lambda self, node: 0} )() cache._clear_pin = lambda node: None cache._record_remove_event = lambda node: None freed = [] cache.cache_controller = type( "Controller", (), { "evict_host": lambda self, indices: freed.append(indices.clone()) or len(indices) }, )() parent = TreeNode() parent.parent = cache.root_node parent.key = RadixKey([1]) parent.value = None parent.host_len = 4 parent.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.tensor([2], dtype=torch.int64), host_indices=torch.tensor([80], dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache.root_node.children[1] = parent stale_child = TreeNode() stale_child.parent = parent stale_child.key = RadixKey([2]) stale_child.value = None stale_child.host_len = 0 stale_child.cp_hicache = None parent.children[2] = stale_child sibling = TreeNode() sibling.parent = parent sibling.key = RadixKey([3]) sibling.value = None sibling.host_len = 0 sibling.cp_hicache = None parent.children[3] = sibling cache.evictable_host_leaves.add(stale_child) physical_freed = cache._evict_host_for_physical_slots(1) self.assertEqual(physical_freed, 0) self.assertEqual(freed, []) self.assertNotIn(2, parent.children) self.assertIn(1, cache.root_node.children) self.assertIs(cache.root_node.children[1], parent) self.assertIn(3, parent.children) self.assertIs(parent.children[3], sibling) self.assertEqual(parent.host_len, 4) self.assertIsNotNone(parent.cp_hicache) self.assertEqual(parent.cp_hicache.host_indices.tolist(), [80]) def test_synchronized_cp_host_eviction_removes_zero_owned_logical_leaf(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 2 cache.tp_group = object() cache._tp_group_rank = 0 cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() cache.get_child_key_fn = lambda key: key.token_ids[0] cache.eviction_strategy = type( "Strategy", (), {"get_priority": lambda self, node: 0} )() cache._clear_pin = lambda node: None cache._record_remove_event = lambda node: None freed = [] cache.cache_controller = type( "Controller", (), { "evict_host": lambda self, indices: freed.append(indices.clone()) or len(indices) }, )() node = TreeNode() node.parent = cache.root_node node.key = RadixKey([1, 2, 3, 4]) node.value = None node.host_len = 4 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache.root_node.children[1] = node cache.evictable_host_leaves.add(node) all_done_states = iter([0, 1]) def fake_all_reduce(done, op=None, group=None): done.fill_(next(all_done_states, 1)) with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce): physical_freed = cache._evict_host_for_physical_slots( 0, synchronize_across_ranks=True ) self.assertEqual(physical_freed, 0) self.assertEqual(freed, []) self.assertNotIn(1, cache.root_node.children) self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) def test_cp_host_eviction_skips_all_reduce_without_tp_group(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 2 cache.tp_group = None cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() cache.get_child_key_fn = lambda key: key.token_ids[0] cache.eviction_strategy = type( "Strategy", (), {"get_priority": lambda self, node: 0} )() cache._clear_pin = lambda node: None cache._record_remove_event = lambda node: None cache.cache_controller = type( "Controller", (), { "evict_host": lambda self, indices: (_ for _ in ()).throw( AssertionError("must not evict host slots") ) }, )() node = TreeNode() node.parent = cache.root_node node.key = RadixKey([1, 2, 3, 4]) node.value = None node.host_len = 4 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) cache.root_node.children[1] = node cache.evictable_host_leaves.add(node) with patch( "torch.distributed.all_reduce", side_effect=AssertionError("must not all_reduce without tp_group"), ): physical_freed = cache._evict_host_for_physical_slots( 0, synchronize_across_ranks=True ) self.assertEqual(physical_freed, 0) self.assertIn(1, cache.root_node.children) self.assertEqual(node.host_len, 4) self.assertIsNotNone(node.cp_hicache) class TestHiRadixCacheCPLoadBack(CustomTestCase): def test_cp_load_back_uses_host_len_not_host_value(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.device = "cpu" cache.load_back_threshold = 1 cache.evictable_size_ = 0 cache.metrics_collector = None cache.ongoing_load_back = {} cache.cache_controller = type( "Controller", (), { "load_cp": lambda self, nodes, node_id=-1: torch.arange( 32, 40, dtype=torch.int64 ) }, )() cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})() cache.dec_lock_ref = lambda node: None cache.evict = lambda params: None parent = cache.root_node parent.key = RadixKey([]) parent.value = torch.empty((0,), dtype=torch.int64) node = TreeNode() node.parent = parent node.key = RadixKey(list(range(8))) node.value = None node.host_value = None node.host_len = 8 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.tensor([0, 1], dtype=torch.int64), host_indices=torch.tensor([50, 51], dtype=torch.int64), page_owners=torch.zeros(max(8, 0), dtype=torch.int8), page_size=1, ) loaded = cache.load_back(node) self.assertEqual(loaded.tolist(), list(range(32, 40))) self.assertEqual(node.value.tolist(), list(range(32, 40))) def test_cp_load_back_threshold_uses_logical_length(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.root_node.value = torch.empty((0,), dtype=torch.int64) cache.load_back_threshold = 5 cache.evictable_size_ = 0 cache.metrics_collector = None cache.ongoing_load_back = {} cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})() cache.dec_lock_ref = lambda node: None cache.cache_controller = type( "Controller", (), { "load_cp": lambda self, nodes, node_id=-1: torch.arange( 10, 16, dtype=torch.int64 ) }, )() node = TreeNode() node.parent = cache.root_node node.value = None node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.zeros(max(6, 0), dtype=torch.int8), page_size=1, ) loaded = cache.load_back(node) self.assertEqual(loaded.tolist(), [10, 11, 12, 13, 14, 15]) def test_cp_match_prefix_counts_logical_host_hit_without_host_value(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 1 cache.get_child_key_fn = lambda key: key.token_ids[0] cache.key_match_fn = lambda child_key, key: sum( 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs ) cache.maybe_bigram_convert = lambda key: (key, None) root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.host_len = 0 cache.root_node = root node = TreeNode() node.parent = root node.key = RadixKey(list(range(8))) node.value = None node.host_value = None node.host_len = 8 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.tensor([0, 1], dtype=torch.int64), host_indices=torch.tensor([50, 51], dtype=torch.int64), page_owners=torch.zeros(max(8, 0), dtype=torch.int8), page_size=1, ) root.children[0] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8))))) self.assertEqual(result.host_hit_length, 8) self.assertIs(result.last_host_node, node) self.assertIs(result.last_device_node, root) def test_cp_match_prefix_does_not_admit_inflight_write_as_host_hit(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 1 cache.ongoing_write_through = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True) cache.get_child_key_fn = lambda key: key.token_ids[0] cache.key_match_fn = lambda child_key, key: sum( 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs ) cache.maybe_bigram_convert = lambda key: (key, None) root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.host_len = 0 cache.root_node = root node = TreeNode() node.id = 126 node.parent = root node.key = RadixKey(list(range(8))) node.value = torch.arange(8, dtype=torch.int64) node.host_value = None node.host_len = 8 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.tensor([0, 1], dtype=torch.int64), host_indices=torch.tensor([50, 51], dtype=torch.int64), draft_host_indices=torch.tensor([150, 151], dtype=torch.int64), page_owners=torch.zeros(max(8, 0), dtype=torch.int8), page_size=1, ) root.children[0] = node cache.ongoing_write_through[node.id] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8))))) self.assertEqual(result.device_indices.tolist(), list(range(8))) self.assertEqual(result.host_hit_length, 0) self.assertIs(result.last_device_node, node) self.assertIs(result.last_host_node, root) def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 4 cache.maybe_bigram_convert = lambda key: (key, None) cache.get_child_key_fn = lambda key: key.token_ids[0] cache.root_node = TreeNode() result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) self.assertEqual(result.device_indices.tolist(), []) self.assertIs(result.last_device_node, cache.root_node) self.assertIs(result.last_host_node, cache.root_node) self.assertEqual(result.host_hit_length, 0) def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = False cache.device = "cpu" cache.disable = False cache.page_size = 1 cache.get_child_key_fn = lambda key: key.token_ids[0] cache.key_match_fn = lambda child_key, key: sum( 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs ) cache.maybe_bigram_convert = lambda key: (key, None) root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.host_value = None cache.root_node = root node = TreeNode() node.parent = root node.key = RadixKey(list(range(4))) node.value = torch.arange(4, dtype=torch.int64) node.host_value = None root.children[0] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(4))))) self.assertEqual(result.host_hit_length, 0) self.assertIs(result.last_host_node, root) self.assertIs(result.last_device_node, node) def test_cp_load_back_non_evicted_node_returns_none_without_loading(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.load_back_threshold = 1 cache.evictable_size_ = 0 cache.metrics_collector = None cache.ongoing_load_back = {} cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})() cache.dec_lock_ref = lambda node: None cache.cache_controller = type( "Controller", (), { "load_cp": lambda self, nodes, node_id=-1: (_ for _ in ()).throw( AssertionError("must not load without evicted nodes") ) }, )() node = TreeNode() node.value = torch.empty((0,), dtype=torch.int64) loaded = cache.load_back(node) self.assertIsNone(loaded) self.assertEqual(cache.ongoing_load_back, {}) if __name__ == "__main__": unittest.main()