import inspect import re import sys import types import unittest from unittest.mock import MagicMock, 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", "moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> 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", "moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> 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.managers.cache_controller import ( HiCacheAck, HiCacheWriteFailure, HiCacheWriteReservation, ) import sglang.srt.mem_cache.common as mem_cache_common import sglang.srt.mem_cache.hiradix_cache as hiradix_cache from sglang.srt.mem_cache.base_prefix_cache import ( EvictParams, InsertParams, MatchPrefixParams, ) from sglang.srt.mem_cache.hiradix_cache import ( CpHiCacheCapacitySnapshot, CpHiCacheEvictionPlan, CpHiCacheNodeMetadata, CpWriteAdmission, HiRadixCache, PendingHiCacheBackup, PreparedCpHiCacheBackup, _compute_shared_hicache_token_capacities, ) from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, _key_match_paged from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache 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 TestHiCacheEvictLoggingLevels(CustomTestCase): def assert_debug_marker(self, source: str, marker: str): self.assertIn(marker, source) self.assertRegex( source, r"logger\.debug\(\s*\n\s*\"" + re.escape(marker), msg=f"{marker} should be debug-only on the success/no-op hot path", ) self.assertNotRegex( source, r"logger\.info\(\s*\n\s*\"" + re.escape(marker), msg=f"{marker} must not stay at INFO on the success/no-op hot path", ) def test_evict_hot_path_success_logs_are_debug_only(self): common_source = inspect.getsource(mem_cache_common) hiradix_source = inspect.getsource(hiradix_cache) for marker in ( "[MemCache-evict] evict_from_tree_cache:", "[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size", "[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d:", "[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result:", ): self.assert_debug_marker(common_source, marker) for marker in ( "[HiCache-write] write_backup triggered: node_id=%d", "[HiCache-write] write_backup CP SUBMITTED: node_id=%d", "[HiCache-write] write_backup non-CP SUCCESS: node_id=%d", "[HiCache-write] write_backup non-CP FAILED (host full): node_id=%d", "[HiCache-write] _inc_hit_count triggering write_through: node_id=%d", "[HiCache-write] writing_check released %d write locks:", "[HiCache-load] load_back CP: node_id=%d", "[HiCache-load] load_back CP SUCCESS: node_id=%d", "[HiCache-load] load_back non-CP SUCCESS: node_id=%d", "[HiCache-collective] tag=%s", ): self.assert_debug_marker(hiradix_source, marker) import sglang.srt.managers.cache_controller as cache_controller_mod cache_controller_source = inspect.getsource(cache_controller_mod) for marker in ( "[CacheCtrl-write] write: node_id=%d", "[CacheCtrl-write] write non-CP submitted: node_id=%d", "[CacheCtrl-write] submit_write_cp_all_layer submitted: node_id=%d", "[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d", ): self.assert_debug_marker(cache_controller_source, marker) self.assertRegex( hiradix_source, r"logger\.warning\(\s*\n\s*\"\[HiCache-write\] write_backup CP FAILED after deterministic retry:", ) self.assertNotRegex( hiradix_source, r"logger\.info\(\s*\n\s*\"\[HiCache-write\] write_backup CP FAILED after deterministic retry:", ) 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_fp8_nsa_hicache_size_estimate_uses_packed_row_width(self): pool = object.__new__(NSATokenToKVPool) pool.store_dtype = torch.uint8 pool.kv_cache_dim = 656 pool.layer_num = 78 pool.index_head_dim = 128 pool.quant_block_size = 128 size_per_token = hiradix_cache._estimate_hicache_size_per_token(pool) # FP8 NSA stores packed MLA rows as 656 uint8 bytes/token/layer plus # the paged index buffer: 128 uint8 K bytes + 4 uint8 scale bytes. self.assertEqual(size_per_token, (656 + 128 + 4) * 78) def test_fp8_nsa_hicache_size_estimate_uses_active_index_layers(self): pool = object.__new__(NSATokenToKVPool) pool.store_dtype = torch.uint8 pool.kv_cache_dim = 656 pool.layer_num = 78 pool.index_head_dim = 128 pool.quant_block_size = 128 pool.index_active_layer_ids = tuple(range(0, 78, 4)) size_per_token = hiradix_cache._estimate_hicache_size_per_token(pool) indexer_size_per_token = 128 + 4 self.assertEqual( size_per_token, 656 * 78 + indexer_size_per_token * len(pool.index_active_layer_ids), ) def test_fp8_shared_budget_matches_target_and_one_layer_draft_capacity(self): bytes_per_layer = 656 + 128 + 4 target_size_per_token = bytes_per_layer * 78 draft_size_per_token = bytes_per_layer target_tokens, draft_tokens = _compute_shared_hicache_token_capacities( total_host_bytes=int(150 * 1e9), target_size_per_token=target_size_per_token, draft_size_per_token=draft_size_per_token, page_size=64, ) self.assertGreaterEqual(draft_tokens, target_tokens) self.assertLessEqual( target_tokens * target_size_per_token + draft_tokens * draft_size_per_token, int(150 * 1e9), ) # The draft pool is much smaller in bytes because GLM-5 EAGLE draft has # one executable layer, but it still has at least target token capacity. self.assertLess( draft_tokens * draft_size_per_token, target_tokens * target_size_per_token // 32, ) 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_paged_key_match_returns_valid_tail_length_not_next_page(self): key = RadixKey(list(range(6))) self.assertEqual(_key_match_paged(key, key, page_size=4), 6) 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_valid_length_can_be_shorter_than_physical_padded_length(self): metadata = CpHiCacheNodeMetadata( logical_len=100, padded_len=128, owned_positions=torch.tensor([0, 63, 100, 127], dtype=torch.int64), host_indices=torch.tensor([10, 11, 12, 13], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=64, ) self.assertEqual(metadata.logical_len, 100) self.assertEqual(metadata.valid_len, 100) self.assertEqual(metadata.padded_len, 128) self.assertEqual(metadata.page_owners.tolist(), [0, 1]) 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 TestCpHiCacheFreeRoom(CustomTestCase): def test_free_room_deficit_does_not_evict_when_required_fits_trigger(self): deficit = hiradix_cache._free_room_deficit( required=64, available=96, capacity=256, page_size=64, target_ratio=0.5, trigger_ratio=0.0, ) self.assertEqual(deficit, 0) def test_free_room_deficit_evicts_to_target_after_trigger(self): deficit = hiradix_cache._free_room_deficit( required=64, available=96, capacity=256, page_size=64, target_ratio=0.5, trigger_ratio=0.25, ) self.assertEqual(deficit, 128) def test_free_room_deficit_trigger_accounts_for_pending_reservation(self): deficit = hiradix_cache._free_room_deficit( required=64, available=96, capacity=256, page_size=64, target_ratio=0.5, trigger_ratio=0.25, ) # available=96 is above trigger_room=64, but after reserving required=64 # only 32 tokens would remain. Trigger checks available - required, # then evicts at least target_room=128 tokens. self.assertEqual(deficit, 128) def test_free_room_deficit_rounds_room_to_page(self): deficit = hiradix_cache._free_room_deficit( required=64, available=0, capacity=100, page_size=64, target_ratio=0.01, trigger_ratio=0.0, ) self.assertEqual(deficit, 64) def test_cp_host_write_admission_uses_trigger_target_room(self): cache = HiRadixCache.__new__(HiRadixCache) cache.page_size = 64 cache.hicache_host_free_room_ratio = 0.5 cache.hicache_host_free_room_trigger_ratio = 0.25 cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( target_capacity=(256, 256), draft_capacity=None, committed_target=(160, 0), committed_draft=(0, 0), pending_target=(0, 0), pending_draft=(0, 0), ) cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( victims=(), planned_freed=tuple(0 for _ in deficit), remaining_deficit=tuple(0 for _ in deficit), ) admission = cache._cp_build_write_admission( torch.arange(64, dtype=torch.int64), node_id=600, phase="unit", ) self.assertEqual(admission.target_available_by_owner, (96, 256)) self.assertEqual(admission.deficit_by_owner, (128, 0)) def test_cp_host_write_admission_does_not_evict_when_trigger_room_fits(self): cache = HiRadixCache.__new__(HiRadixCache) cache.page_size = 64 cache.hicache_host_free_room_ratio = 0.5 cache.hicache_host_free_room_trigger_ratio = 0.0 cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( target_capacity=(256, 256), draft_capacity=None, committed_target=(160, 0), committed_draft=(0, 0), pending_target=(0, 0), pending_draft=(0, 0), ) cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( victims=(), planned_freed=tuple(0 for _ in deficit), remaining_deficit=tuple(0 for _ in deficit), ) admission = cache._cp_build_write_admission( torch.arange(64, dtype=torch.int64), node_id=601, phase="unit", ) self.assertEqual(admission.target_available_by_owner, (96, 256)) self.assertEqual(admission.deficit_by_owner, (0, 0)) def test_cp_host_write_admission_uses_draft_room_deficit(self): cache = HiRadixCache.__new__(HiRadixCache) cache.page_size = 64 cache.hicache_host_free_room_ratio = 0.5 cache.hicache_host_free_room_trigger_ratio = 0.25 cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0) cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( target_capacity=(256, 256), draft_capacity=(256, 256), committed_target=(0, 0), committed_draft=(192, 0), pending_target=(0, 0), pending_draft=(0, 0), ) cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( victims=(), planned_freed=tuple(0 for _ in deficit), remaining_deficit=tuple(0 for _ in deficit), ) admission = cache._cp_build_write_admission( torch.arange(64, dtype=torch.int64), node_id=602, phase="unit", ) self.assertEqual(admission.target_available_by_owner, (256, 256)) self.assertEqual(admission.draft_available_by_owner, (64, 256)) self.assertEqual(admission.deficit_by_owner, (128, 0)) def test_cp_host_write_batch_admission_uses_aggregate_required(self): cache = HiRadixCache.__new__(HiRadixCache) cache.page_size = 64 cache.hicache_host_free_room_ratio = 0.5 cache.hicache_host_free_room_trigger_ratio = 0.25 cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot( target_capacity=(256, 256), draft_capacity=None, committed_target=(160, 0), committed_draft=(0, 0), pending_target=(0, 0), pending_draft=(0, 0), ) cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan( victims=(), planned_freed=tuple(0 for _ in deficit), remaining_deficit=tuple(0 for _ in deficit), ) admission = cache._cp_build_write_admission_from_required( (128, 0), node_id=603, phase="batch_unit", ) self.assertEqual(admission.required_by_owner, (128, 0)) self.assertEqual(admission.target_available_by_owner, (96, 256)) # required=128, target_room=128, available=96. self.assertEqual(admission.deficit_by_owner, (128, 0)) def test_prepare_write_backups_for_reqs_runs_one_batch_admission(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 1 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(16, dtype=torch.int64).view(2, 8) ) cache.cache_controller = FakeReserveWriteController( [ lambda device_indices, node_id: make_write_reservation( device_indices, node_id=node_id, host_start=700 ), lambda device_indices, node_id: make_write_reservation( device_indices, node_id=node_id, host_start=800 ), ] ) cache._probe_existing_radix_prefix_len_no_split = lambda key: 0 cache._cp_zero_counts = lambda cp_size=None: (0,) cache._cp_add_counts = lambda lhs, rhs: (lhs[0] + rhs[0],) cache._cp_required_host_tokens_by_rank = lambda indices: (int(indices.numel()),) batch_admissions = [] evicted_admissions = [] def build_batch_admission(required, *, node_id, phase): batch_admissions.append((required, node_id, phase)) return CpWriteAdmission( node_id=node_id, phase=phase, required_by_owner=required, target_available_by_owner=(0,), draft_available_by_owner=(0,), deficit_by_owner=(1,), eviction_plan=CpHiCacheEvictionPlan( victims=(), planned_freed=(1,), remaining_deficit=(0,), ), ) cache._cp_build_write_admission_from_required = build_batch_admission cache._evict_cp_host_for_write_admission = ( lambda admission, **_: evicted_admissions.append(admission) or True ) req0 = types.SimpleNamespace( rid="batch-rid-0", fill_ids=list(range(4)), cache_protected_len=0, req_pool_idx=0, is_chunked=0, cp_hicache_prepared_backup=None, ) req1 = types.SimpleNamespace( rid="batch-rid-1", fill_ids=list(range(6)), cache_protected_len=0, req_pool_idx=1, is_chunked=0, cp_hicache_prepared_backup=None, ) cache.prepare_write_backups_for_reqs([req0, req1]) self.assertEqual(len(batch_admissions), 1) self.assertEqual(batch_admissions[0][0], (10,)) self.assertEqual(batch_admissions[0][2], "batch_prepare") self.assertEqual(len(evicted_admissions), 1) self.assertIsNotNone(req0.cp_hicache_prepared_backup) self.assertIsNotNone(req1.cp_hicache_prepared_backup) self.assertEqual(len(cache.cache_controller.submitted), 2) 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 FakeReserveWriteController: write_policy = "write_through" has_draft_hicache = False def __init__(self, results): self.results = list(results) self.reservations = [] self.submitted = [] self.submit_kwargs = [] self.evicted_host_indices = [] def reserve_write_cp(self, device_indices, priority=None, node_id=-1): self.reservations.append((device_indices.clone(), node_id)) result = self.results.pop(0) if not callable(result): return result try: return result(device_indices, node_id=node_id) except TypeError: return result(device_indices) def submit_write_cp_all_layer(self, reservation): self.submitted.append(reservation) def submit_write_cp_per_layer(self, reservation, **kwargs): self.submitted.append(reservation) self.submit_kwargs.append(kwargs) def evict_cp_host(self, metadata): self.evicted_host_indices.append(metadata.host_indices.clone()) return len(metadata.host_indices) def make_write_reservation(device_indices, node_id=0, host_start=90): host_indices = torch.arange(host_start, host_start + len(device_indices)) metadata = CpHiCacheNodeMetadata( logical_len=len(device_indices), owned_positions=torch.arange(len(device_indices), dtype=torch.int64), host_indices=host_indices, page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8), page_size=1, ) return HiCacheWriteReservation( metadata=metadata, host_indices=host_indices, physical_device_indices=device_indices.clone(), node_id=node_id, ) class FakeEvictionStrategy: def get_priority(self, node): return 0 class FakeCpLayout: def __init__(self, cp_size=4, cp_rank=0, page_size=1): self.cp_size = cp_size self.cp_rank = cp_rank self.page_size = page_size def owner_for_logical_pages(self, logical_pages): owners = torch.remainder(logical_pages - 1, self.cp_size) return torch.where(logical_pages <= 0, torch.full_like(owners, -1), owners) def logical_locs_to_physical(self, logical_locs): return logical_locs def owned_by_this_rank(self, logical_locs): logical_pages = torch.div( logical_locs, self.page_size, rounding_mode="floor" ) return self.owner_for_logical_pages(logical_pages) == self.cp_rank class FakeEvictDeviceController: write_policy = "write_through" def evict_device(self, device_indices): return len(device_indices) class FakeTokenAllocator: def available_size(self): return 0 def compute_owner_lane_stats(self, page_owners): if len(page_owners) == 0: return [], [], [] cp_size = max(int(owner) for owner in page_owners) + 1 required = [0] * cp_size for owner in page_owners: required[int(owner)] += 1 available = list(required) deficits = [0] * cp_size return required, available, deficits def allocator_state_str(self): return "FakeTokenAllocator" class RecordingTokenAllocator: def __init__(self): self.freed = [] def free(self, free_index): if free_index.numel() > 0: self.freed.append(free_index.detach().cpu().tolist()) class RecordingDeviceAllocator: def __init__(self): self.freed = [] def free(self, free_index): self.freed.append(free_index.detach().cpu().tolist()) class TestHiRadixCacheCPBackup(CustomTestCase): def _minimal_cp_hiradix_cache(self, *, page_size=64): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.page_size = page_size cache.is_eagle = False cache.enable_storage = False cache.enable_kv_cache_events = False cache.evictable_size_ = 0 cache.protected_size_ = 0 cache.evictable_leaves = set() cache.evictable_host_leaves = set() cache.get_child_key_fn = lambda key: tuple(key.token_ids[:page_size]) cache.key_match_fn = lambda key0, key1: _key_match_paged( key0, key1, page_size ) cache.cache_controller = types.SimpleNamespace(write_policy="write_back") cache._record_store_event = lambda node: None cache._record_remove_event = lambda node: None root = TreeNode() root.key = RadixKey(token_ids=[], extra_key=None) root.value = torch.empty((0,), dtype=torch.int64) root.children = {} root.parent = None cache.root_node = root return cache def test_cp_eagle_finished_cache_preserves_retained_tail_page(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable_finished_insert = False cache.disable = False cache.is_eagle = True cache.page_size = 64 cache._uses_cp_hicache = True cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(128, dtype=torch.int64).view(1, 128) ) allocator = RecordingTokenAllocator() cache.token_to_kv_pool_allocator = allocator cache.insert = lambda params: types.SimpleNamespace(prefix_len=0) cache.dec_lock_ref = lambda node: None req = types.SimpleNamespace( origin_input_ids=[10, 11, 12, 13], output_ids=[], req_pool_idx=0, extra_key=None, cache_protected_len=0, last_node=object(), cp_hicache_prepared_backup=None, pop_committed_kv_cache=lambda: 4, ) cache.cache_finished_req(req) self.assertEqual(allocator.freed, []) def test_cp_valid_tail_device_accounting_uses_physical_page_span(self): cache = self._minimal_cp_hiradix_cache(page_size=64) result = cache.insert( InsertParams( key=RadixKey(token_ids=[1, 2, 3], extra_key=None), value=torch.arange(3, dtype=torch.int64), ) ) self.assertEqual(result.prefix_len, 0) self.assertEqual(cache.evictable_size_, 64) node = next(iter(cache.root_node.children.values())) cache.inc_node_lock_ref(node) self.assertEqual(cache.evictable_size_, 0) self.assertEqual(cache.protected_size_, 64) cache.dec_node_lock_ref(node) self.assertEqual(cache.evictable_size_, 64) self.assertEqual(cache.protected_size_, 0) def test_cp_valid_tail_regular_evict_reports_and_subtracts_physical_page(self): cache = self._minimal_cp_hiradix_cache(page_size=64) device_allocator = RecordingDeviceAllocator() cache.cache_controller.mem_pool_device_allocator = device_allocator cache.insert( InsertParams( key=RadixKey(token_ids=[1, 2, 3], extra_key=None), value=torch.arange(3, dtype=torch.int64), ) ) node = next(iter(cache.root_node.children.values())) num_evicted = cache._evict_regular(node) self.assertEqual(num_evicted, 64) self.assertEqual(cache.evictable_size_, 0) self.assertEqual(device_allocator.freed, [[0, 1, 2]]) self.assertEqual(cache.root_node.children, {}) def test_cp_partial_split_floors_unbacked_valid_tail_to_page_boundary(self): cache = self._minimal_cp_hiradix_cache(page_size=64) node = TreeNode() node.host_len = 0 node.cp_hicache = None self.assertEqual(cache._cp_floor_backed_partial_split_len(node, 3), 0) self.assertEqual(cache._cp_floor_backed_partial_split_len(node, 70), 64) self.assertEqual(cache._cp_floor_backed_partial_split_len(node, 128), 128) def test_session_aware_cache_forwards_cp_hicache_prepare(self): calls = [] class Inner: def prepare_write_backup_for_req(self, req): calls.append(req) wrapper = SessionAwareCache(Inner()) req = object() wrapper.prepare_write_backup_for_req(req) self.assertEqual(calls, [req]) def test_session_aware_cache_forwards_cp_hicache_batch_prepare(self): calls = [] class Inner: def prepare_write_backups_for_reqs(self, reqs): calls.append(list(reqs)) wrapper = SessionAwareCache(Inner()) reqs = [object(), object()] wrapper.prepare_write_backups_for_reqs(reqs) self.assertEqual(calls, [reqs]) 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_node_backuped_excludes_explicit_pending_cp_backup_until_commit(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.ongoing_write_through = {} cache.pending_host_backups = {} dec_locked = [] cache.dec_node_lock_ref = lambda node: dec_locked.append(node) node = TreeNode() node.id = 128 node.host_len = 64 metadata = 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, ) node.cp_hicache = metadata cache.pending_host_backups[node.id] = PendingHiCacheBackup( node=node, metadata=metadata, logical_len=64 ) self.assertFalse(cache._node_backuped(node)) cache._commit_pending_backup(node.id) self.assertTrue(cache._node_backuped(node)) self.assertEqual(dec_locked, [node]) def test_rollback_pending_cp_backup_frees_reserved_host_slots(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True evicted = [] cache.cache_controller = types.SimpleNamespace( evict_cp_host=lambda metadata: evicted.append(metadata) or 2, cancel_layer_write_state=lambda node_id: False, ack_write_queue=[], ) dec_locked = [] cache.dec_node_lock_ref = lambda node: dec_locked.append(node) cache.pending_host_backups = {} node = TreeNode() node.id = 129 metadata = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.tensor([0, 1], dtype=torch.int64), host_indices=torch.tensor([10, 11], dtype=torch.int64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) cache.pending_host_backups[node.id] = PendingHiCacheBackup( node=node, metadata=metadata, logical_len=64 ) cache._rollback_pending_backup(node.id) self.assertEqual(evicted, [metadata]) self.assertEqual(dec_locked, [node]) self.assertNotIn(node.id, cache.pending_host_backups) 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_uses_deterministic_host_eviction_before_reserve(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.page_size = 1 reservation_factory = lambda device_indices: make_write_reservation( device_indices, node_id=122 ) cache.cache_controller = FakeReserveWriteController([reservation_factory]) cache.cache_controller.cp_shared_kv_layout = FakeCpLayout(cp_size=1, cp_rank=0) cache.token_to_kv_pool_host = types.SimpleNamespace(size=16) 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._update_host_leaf_status = lambda node: None cache.ongoing_write_through = {} cache.pending_host_backups = {} 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.arange(4, dtype=torch.int64), host_indices=torch.arange(55, 59, 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.id = 122 node.value = torch.arange(16, dtype=torch.int64) cache.write_backup(node) self.assertEqual( cache.cache_controller.evicted_host_indices[0].tolist(), list(range(55, 59)), ) self.assertEqual(evictable_node.host_len, 0) self.assertIsNone(evictable_node.cp_hicache) self.assertNotIn(1, root.children) self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) self.assertIn(node.id, cache.pending_host_backups) self.assertEqual(len(cache.cache_controller.submitted), 1) def test_write_backup_cp_post_forward_path_warns_fallback(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = FakeReserveWriteController( [ lambda device_indices, node_id: make_write_reservation( device_indices, node_id=node_id ) ] ) cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.inc_node_lock_ref = lambda node: None node = TreeNode() node.id = 133 node.value = torch.arange(16, dtype=torch.int64) with self.assertLogs( "sglang.srt.mem_cache.hiradix_cache", level="WARNING" ) as captured: cache.write_backup(node) self.assertTrue( any( "[CP_HICACHE_FALLBACK][post_forward_catch_up_backup]" in message for message in captured.output ) ) def test_prepare_write_backup_for_req_chunked_reserves_completed_pages(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.cache_controller = FakeReserveWriteController([]) cache.is_eagle = False cache.page_size = 4 cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.root_node.children = {} cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) ) cache._reserve_write_cp_indices_no_collective = ( lambda indices, node_id: make_write_reservation( indices, node_id=node_id, host_start=300 ) ) req = types.SimpleNamespace( rid="rid-chunked", is_chunked=1, fill_ids=list(range(10)), req_pool_idx=0, extra_key=None, cache_protected_len=0, cp_hicache_prepared_backup=None, ) cache.prepare_write_backup_for_req(req) self.assertEqual(len(cache.cache_controller.submitted), 1) self.assertEqual( cache.cache_controller.submitted[0].physical_device_indices.tolist(), list(range(8)), ) self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 8) def test_cp_cache_unfinished_chunked_inserts_only_completed_pages(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 class Pool: def __init__(self): self.req_to_token = torch.arange(16, dtype=torch.int64).view(1, 16) self.writes = [] def write(self, index, values): self.writes.append((index, values.detach().cpu().tolist())) pool = Pool() cache.req_to_token_pool = pool cache.token_to_kv_pool_allocator = RecordingTokenAllocator() inserted = [] def insert(params): inserted.append(params) return types.SimpleNamespace(prefix_len=8, pending_backup_deferred_node=None) cache.insert = insert cache.match_prefix = lambda params: types.SimpleNamespace( device_indices=torch.arange(8, dtype=torch.int64), last_device_node=TreeNode(), ) cache._free_kv_indices_range = lambda *args, **kwargs: None cache.dec_lock_ref = lambda node: None cache.inc_lock_ref = lambda node: None req = types.SimpleNamespace( fill_ids=list(range(10)), req_pool_idx=0, extra_key=None, cache_protected_len=0, last_node=TreeNode(), prefix_indices=torch.empty((0,), dtype=torch.int64), cp_hicache_prepared_backup=None, ) cache.cache_unfinished_req(req, chunked=True) self.assertEqual(len(inserted), 1) self.assertEqual(inserted[0].key.token_ids, list(range(8))) self.assertEqual(inserted[0].value.tolist(), list(range(8))) self.assertEqual( pool.writes, [((0, slice(0, 8, None)), list(range(8)))], ) self.assertEqual(req.cache_protected_len, 8) self.assertEqual(req.prefix_indices.tolist(), list(range(10))) def test_write_backup_deterministic_eviction_avoids_reserve_all_reduce(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.page_size = 1 cache.tp_world_size = 2 cache.tp_group = object() cache.cache_controller = FakeReserveWriteController( [ lambda device_indices: make_write_reservation( device_indices, node_id=130, host_start=100 ), ] ) cache.cache_controller.cp_shared_kv_layout = FakeCpLayout(cp_size=1, cp_rank=0) cache.token_to_kv_pool_host = types.SimpleNamespace(size=16) evictions = [] cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: ( evictions.append((required, synchronize_across_ranks)) or required ) 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._update_host_leaf_status = lambda node: None cache.ongoing_write_through = {} cache.pending_host_backups = {} 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=[2], extra_key=None) evictable_node.value = None evictable_node.host_len = 16 evictable_node.cp_hicache = CpHiCacheNodeMetadata( logical_len=16, owned_positions=torch.arange(16, dtype=torch.int64), host_indices=torch.arange(90, 106, dtype=torch.int64), page_owners=torch.zeros(16, dtype=torch.int8), page_size=1, ) root.children[2] = evictable_node cache.evictable_host_leaves.add(evictable_node) node = TreeNode() node.id = 130 node.value = torch.arange(16, dtype=torch.int64) with patch( "torch.distributed.all_reduce", side_effect=AssertionError("reserve admission must not all_reduce"), ): backed_len = cache.write_backup(node) self.assertEqual(backed_len, 16) self.assertEqual( cache.cache_controller.evicted_host_indices[0].tolist(), list(range(90, 106)), ) self.assertEqual(evictions, []) self.assertEqual(len(cache.cache_controller.submitted), 1) self.assertEqual( cache.cache_controller.submitted[0].metadata.host_indices.tolist(), list(range(100, 116)), ) def test_insert_attaches_prepared_cp_backup_without_catchup_copy(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_size_ = 0 cache.get_child_key_fn = lambda key: key.token_ids[0] cache.key_match_fn = lambda lhs, rhs: 0 cache.maybe_bigram_convert = lambda key, value: (key, value) cache.is_eagle = False cache.enable_storage = False cache.enable_kv_cache_events = False cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache._record_store_event = lambda node: None cache.ongoing_write_through = {} cache.pending_host_backups = {} locked_nodes = [] cache.inc_node_lock_ref = lambda node: locked_nodes.append(node) cache.cache_controller = types.SimpleNamespace(write_policy="write_through") value = torch.arange(16, dtype=torch.int64) reservation = make_write_reservation(value, node_id=131, host_start=110) prepared = PreparedCpHiCacheBackup( node_id=131, reservation=reservation, metadata=reservation.metadata, logical_len=16, ) result = cache.insert( InsertParams( key=RadixKey(list(range(16))), value=value, cp_hicache_prepared_backup=prepared, ) ) node = cache.root_node.children[0] self.assertEqual(result.prefix_len, 0) self.assertEqual(node.id, 131) self.assertTrue(prepared.attached) self.assertIs(cache.ongoing_write_through[131], node) self.assertIs(cache.pending_host_backups[131].node, node) self.assertIs(cache.pending_host_backups[131].metadata, reservation.metadata) self.assertEqual(cache.pending_host_backups[131].logical_len, 16) self.assertEqual(locked_nodes, [node]) def test_prepare_write_backup_for_req_registers_before_forward_without_catchup( self, ): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 1 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) ) cache.cache_controller = FakeReserveWriteController( [ lambda device_indices, node_id: make_write_reservation( device_indices, node_id=node_id, host_start=130 ) ] ) req = types.SimpleNamespace( rid="rid-prepare", fill_ids=list(range(16)), cache_protected_len=0, req_pool_idx=0, is_chunked=0, cp_hicache_prepared_backup=None, ) cache.prepare_write_backup_for_req(req) prepared = req.cp_hicache_prepared_backup self.assertIsNotNone(prepared) self.assertEqual(prepared.logical_len, 16) self.assertEqual(len(cache.cache_controller.submitted), 1) self.assertIs(cache.cache_controller.submitted[0], prepared.reservation) self.assertEqual( cache.cache_controller.submit_kwargs, [{"catch_up_all_layers": False}], ) def test_prepare_write_backup_for_req_keeps_valid_tail_length(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8) ) cache.cache_controller = FakeReserveWriteController( [ lambda device_indices, node_id: make_write_reservation( device_indices, node_id=node_id, host_start=170 ) ] ) cache.maybe_bigram_convert = lambda key: (key, None) cache.root_node = TreeNode() cache.root_node.children = {} cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) req = types.SimpleNamespace( rid="rid-tail-prepare", fill_ids=list(range(6)), cache_protected_len=0, req_pool_idx=0, is_chunked=0, cp_hicache_prepared_backup=None, ) cache.prepare_write_backup_for_req(req) self.assertIsNotNone(req.cp_hicache_prepared_backup) self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 6) self.assertEqual(cache.cache_controller.reservations[0][0].tolist(), list(range(6))) def test_prepare_write_backup_for_req_floors_mid_page_prefix_hit(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) ) cache.cache_controller = FakeReserveWriteController( [ lambda device_indices, node_id: make_write_reservation( device_indices, node_id=node_id, host_start=180 ) ] ) cache._probe_existing_radix_prefix_len_no_split = lambda key: 6 req = types.SimpleNamespace( rid="rid-tail-extend", fill_ids=list(range(10)), cache_protected_len=6, req_pool_idx=0, is_chunked=0, cp_hicache_prepared_backup=None, ) cache.prepare_write_backup_for_req(req) self.assertIsNotNone(req.cp_hicache_prepared_backup) self.assertEqual( cache.cache_controller.reservations[0][0].tolist(), list(range(4, 10)), ) def test_cache_finished_req_keeps_cp_valid_tail_insert_key(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable_finished_insert = False cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8) ) freed = [] cache.token_to_kv_pool_allocator = types.SimpleNamespace( free=lambda indices: freed.append(indices.clone()) ) inserted = [] cache.insert = lambda params: inserted.append(params) or types.SimpleNamespace( prefix_len=0 ) cache.dec_lock_ref = lambda node: None req = types.SimpleNamespace( origin_input_ids=list(range(6)), output_ids=[], req_pool_idx=0, extra_key=None, cache_protected_len=0, last_node=TreeNode(), cp_hicache_prepared_backup=None, pop_committed_kv_cache=lambda: 6, ) cache.cache_finished_req(req) self.assertEqual(len(inserted), 1) self.assertEqual(inserted[0].key.token_ids, list(range(6))) self.assertEqual(inserted[0].value.tolist(), list(range(6))) self.assertEqual([indices.tolist() for indices in freed], [[], []]) def test_cache_finished_req_cp_insert_duplicate_free_skips_partial_page(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable_finished_insert = False cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8) ) allocator = RecordingTokenAllocator() cache.token_to_kv_pool_allocator = allocator cache.insert = lambda params: types.SimpleNamespace(prefix_len=3) cache.dec_lock_ref = lambda node: None req = types.SimpleNamespace( origin_input_ids=list(range(6)), output_ids=[], req_pool_idx=0, extra_key=None, cache_protected_len=1, last_node=TreeNode(), cp_hicache_prepared_backup=None, pop_committed_kv_cache=lambda: 6, ) cache.cache_finished_req(req) self.assertEqual(allocator.freed, []) def test_cache_finished_req_cp_no_insert_frees_only_full_unprotected_pages(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable_finished_insert = False cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(12, dtype=torch.int64).view(1, 12) ) allocator = RecordingTokenAllocator() cache.token_to_kv_pool_allocator = allocator cache.dec_lock_ref = lambda node: None req = types.SimpleNamespace( origin_input_ids=list(range(9)), output_ids=[], req_pool_idx=0, extra_key=None, cache_protected_len=1, last_node=TreeNode(), cp_hicache_prepared_backup=None, pop_committed_kv_cache=lambda: 9, ) cache.cache_finished_req(req, is_insert=False) self.assertEqual(allocator.freed, [[4, 5, 6, 7, 8]]) def test_cache_unfinished_req_cp_duplicate_free_skips_partial_page(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 writes = [] class Pool: req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8) def write(self, index, values): writes.append((index, values.clone())) cache.req_to_token_pool = Pool() allocator = RecordingTokenAllocator() cache.token_to_kv_pool_allocator = allocator cache.insert = lambda params: types.SimpleNamespace(prefix_len=3) new_last_node = TreeNode() cache.match_prefix = lambda params: types.SimpleNamespace( device_indices=torch.arange(6, dtype=torch.int64), last_device_node=new_last_node, ) cache.dec_lock_ref = lambda node: None cache.inc_lock_ref = lambda node: None req = types.SimpleNamespace( fill_ids=list(range(6)), req_pool_idx=0, extra_key=None, cache_protected_len=1, last_node=TreeNode(), prefix_indices=torch.empty((0,), dtype=torch.int64), cp_hicache_prepared_backup=None, ) cache.cache_unfinished_req(req) self.assertEqual(allocator.freed, []) self.assertEqual(req.cache_protected_len, 6) def test_cache_unfinished_req_deferred_insert_preserves_request_kv(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 deferred_node = TreeNode() class Pool: req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8) def write(self, index, values): raise AssertionError("deferred insert must not rewrite req pool") cache.req_to_token_pool = Pool() cache.token_to_kv_pool_allocator = RecordingTokenAllocator() cache.insert = lambda params: types.SimpleNamespace( prefix_len=1, pending_backup_deferred_node=deferred_node, ) cache.match_prefix = lambda params: (_ for _ in ()).throw( AssertionError("deferred insert must not rematch") ) cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("deferred unfinished insert must keep request KV") ) old_last_node = TreeNode() req = types.SimpleNamespace( fill_ids=list(range(6)), req_pool_idx=0, extra_key=None, cache_protected_len=1, last_node=old_last_node, prefix_indices=torch.empty((0,), dtype=torch.int64), cp_hicache_prepared_backup=None, ) cache.cache_unfinished_req(req, chunked=True) self.assertEqual(req.cache_protected_len, 1) self.assertIs(req.last_node, old_last_node) self.assertEqual(req.prefix_indices.tolist(), list(range(6))) self.assertEqual(cache.token_to_kv_pool_allocator.freed, []) def test_cache_unfinished_req_deferred_insert_preserves_unattached_backup_for_final_insert( self, ): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 deferred_node = TreeNode() rollbacks = [] class Pool: req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8) cache.req_to_token_pool = Pool() cache.token_to_kv_pool_allocator = RecordingTokenAllocator() cache.insert = lambda params: types.SimpleNamespace( prefix_len=1, pending_backup_deferred_node=deferred_node, ) cache.match_prefix = lambda params: (_ for _ in ()).throw( AssertionError("deferred insert must not rematch") ) cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("deferred unfinished insert must keep request KV") ) cache._rollback_prepared_cp_backup = lambda prepared, reason: rollbacks.append( (prepared, reason) ) prepared = types.SimpleNamespace(attached=False) req = types.SimpleNamespace( fill_ids=list(range(6)), req_pool_idx=0, extra_key=None, cache_protected_len=1, last_node=TreeNode(), prefix_indices=torch.empty((0,), dtype=torch.int64), cp_hicache_prepared_backup=prepared, ) cache.cache_unfinished_req(req, chunked=False) self.assertEqual(rollbacks, []) self.assertIs(req.cp_hicache_prepared_backup, prepared) self.assertEqual(req.prefix_indices.tolist(), list(range(6))) def test_cache_unfinished_chunked_deferred_insert_rolls_back_unattached_backup( self, ): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 4 deferred_node = TreeNode() rollbacks = [] class Pool: req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8) cache.req_to_token_pool = Pool() cache.token_to_kv_pool_allocator = RecordingTokenAllocator() cache.insert = lambda params: types.SimpleNamespace( prefix_len=1, pending_backup_deferred_node=deferred_node, ) cache.match_prefix = lambda params: (_ for _ in ()).throw( AssertionError("deferred insert must not rematch") ) cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("deferred unfinished insert must keep request KV") ) cache._rollback_prepared_cp_backup = lambda prepared, reason: rollbacks.append( (prepared, reason) ) prepared = types.SimpleNamespace(attached=False) req = types.SimpleNamespace( fill_ids=list(range(6)), req_pool_idx=0, extra_key=None, cache_protected_len=1, last_node=TreeNode(), prefix_indices=torch.empty((0,), dtype=torch.int64), cp_hicache_prepared_backup=prepared, ) cache.cache_unfinished_req(req, chunked=True) self.assertEqual( rollbacks, [(prepared, "unfinished_pending_backup_split_deferred")] ) self.assertIsNone(req.cp_hicache_prepared_backup) self.assertEqual(req.prefix_indices.tolist(), list(range(6))) def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.is_eagle = False cache.page_size = 1 cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.root_node.children = {} cache.get_child_key_fn = lambda key: key.token_ids[0] def key_match(lhs, rhs): matched = 0 for left, right in zip(lhs.token_ids, rhs.token_ids): if left != right: break matched += 1 return matched cache.key_match_fn = key_match cache.maybe_bigram_convert = lambda key, value=None: (key, value) cache.req_to_token_pool = types.SimpleNamespace( req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) ) cache.cache_controller = FakeReserveWriteController([]) existing = TreeNode() existing.id = 201 existing.parent = cache.root_node existing.key = RadixKey(list(range(16))) existing.value = torch.arange(16, dtype=torch.int64) cache.root_node.children[0] = existing req = types.SimpleNamespace( rid="rid-existing-prefix", fill_ids=list(range(16)), extra_key=None, cache_protected_len=8, req_pool_idx=0, is_chunked=0, cp_hicache_prepared_backup=None, ) cache.prepare_write_backup_for_req(req) self.assertIsNone(req.cp_hicache_prepared_backup) self.assertEqual(cache.cache_controller.reservations, []) self.assertEqual(cache.cache_controller.submitted, []) def test_rollback_unattached_prepared_cp_backup_removes_orphan_ack(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True value = torch.arange(4, dtype=torch.int64) reservation = make_write_reservation(value, node_id=132, host_start=120) prepared = PreparedCpHiCacheBackup( node_id=132, reservation=reservation, metadata=reservation.metadata, logical_len=4, ) evicted = [] class ReadyEvent: def synchronize(self): pass cache.cache_controller = types.SimpleNamespace( pending_layer_writes={}, ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [132])], evict_cp_host=lambda metadata: evicted.append(metadata), ) cache._rollback_prepared_cp_backup(prepared, "test") self.assertEqual(cache.cache_controller.ack_write_queue, []) self.assertEqual(evicted, [reservation.metadata]) def test_writing_check_defers_unattached_prepared_ack(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 1 cache.enable_storage = False cache.pending_host_backups = {} class ReadyEvent: def query(self): return True def synchronize(self): pass attached_node = TreeNode() attached_node.id = 53 dec_locked = [] cache.ongoing_write_through = {53: attached_node} cache.dec_node_lock_ref = lambda node: dec_locked.append(node.id) cache.cache_controller = types.SimpleNamespace( ack_write_queue=[ HiCacheAck(ReadyEvent(), ReadyEvent(), [54]), HiCacheAck(ReadyEvent(), ReadyEvent(), [53]), ], ) cache.writing_check() self.assertEqual( [ack.node_ids for ack in cache.cache_controller.ack_write_queue], [[54], [53]], ) self.assertEqual(cache.ongoing_write_through, {53: attached_node}) self.assertEqual(dec_locked, []) @staticmethod def _make_exploding_queue(): class _ExplodingQueue(list): def __iter__(self): raise AssertionError("fresh node ids must not scan the ack queue") def __bool__(self): return True return _ExplodingQueue() def test_reserve_write_cp_refuses_node_with_undrained_ack(self): """Producer-side invariant: at most one in-flight ack per node_id. A reservation for a node whose previous final ack is still queued is refused on the normal HiCacheWriteFailure path instead of creating a second layer-write state (which would enqueue a duplicate ack).""" from sglang.srt.managers.cache_controller import ( HiCacheController, HiCacheWriteFailure, ) controller = types.SimpleNamespace( ack_write_queue=[ HiCacheAck(MagicMock(), MagicMock(), [4614]), ], _max_write_ack_node_id=4614, ) controller.node_has_undrained_write_ack = types.MethodType( HiCacheController.node_has_undrained_write_ack, controller ) result = HiCacheController.reserve_write_cp( controller, device_indices=torch.arange(4, dtype=torch.int64), node_id=4614, ) self.assertIsInstance(result, HiCacheWriteFailure) self.assertEqual(result.required_host_slots, 0) # An old (<= max) but absent node id is resolved by the queue scan. self.assertFalse( HiCacheController.node_has_undrained_write_ack(controller, 4613) ) # Fresh ids (> max ever acked, i.e. every prepare-path reservation) # take the O(1) fast path: no queue scan at all. controller.ack_write_queue = self._make_exploding_queue() self.assertFalse( HiCacheController.node_has_undrained_write_ack(controller, 9999) ) def test_reserve_wrapper_skips_capacity_retry_on_undrained_ack(self): """An undrained_ack refusal must return through _reserve_write_cp_indices_no_collective as a skip-this-round failure: entering the capacity retry would either trip the predicted-no-deficit RuntimeError (free host) or trigger pointless evictions (full host).""" from sglang.srt.managers.cache_controller import HiCacheWriteFailure cache = HiRadixCache.__new__(HiRadixCache) refusal = HiCacheWriteFailure(required_host_slots=0, reason="undrained_ack") cache.cache_controller = types.SimpleNamespace( reserve_write_cp=lambda device_indices, node_id: refusal, ) def _admission(device_indices, node_id, phase): if phase != "initial": raise AssertionError( f"undrained_ack must not reach admission phase {phase!r}" ) return types.SimpleNamespace( deficit_by_owner=[0], eviction_plan=types.SimpleNamespace(remaining_deficit=[0]), ) cache._cp_build_write_admission = _admission cache._evict_cp_host_for_write_admission = lambda *a, **k: (_ for _ in ()).throw( AssertionError("undrained_ack must not trigger host eviction") ) result = cache._reserve_write_cp_indices_no_collective( torch.arange(4, dtype=torch.int64), 4614 ) self.assertIs(result, refusal) self.assertEqual(result.reason, "undrained_ack") def test_rollback_pending_backup_scrubs_acks_and_cancels_state(self): """The write_backup exception rollback must not orphan an ack or a half-submitted layer-write state: an orphaned ack re-opens the registration guard while still queued, producing the duplicate-ack crash; a live state would keep writing into evicted host slots.""" cache = HiRadixCache.__new__(HiRadixCache) node = TreeNode() node.id = 4614 metadata = object() dec_locked = [] cache.pending_host_backups = { 4614: PendingHiCacheBackup( node=node, metadata=metadata, logical_len=8, submitted=True, locked=True, ) } cache.dec_node_lock_ref = lambda n: dec_locked.append(n.id) finish_event = MagicMock() cancelled = [] evicted = [] cache.cache_controller = types.SimpleNamespace( ack_write_queue=[ HiCacheAck(MagicMock(), finish_event, [4614]), HiCacheAck(MagicMock(), MagicMock(), [4615]), ], cancel_layer_write_state=lambda node_id: cancelled.append(node_id), evict_cp_host=lambda meta: evicted.append(meta), ) node.cp_hicache = metadata returned = cache._rollback_pending_backup(4614) self.assertIs(returned, node) self.assertEqual(cache.pending_host_backups, {}) self.assertEqual(cancelled, [4614]) self.assertEqual(evicted, [metadata]) # The 4614 ack is scrubbed (after syncing its finish event); the # unrelated 4615 ack is retained. finish_event.synchronize.assert_called_once() self.assertEqual( [ack.node_ids for ack in cache.cache_controller.ack_write_queue], [[4615]], ) self.assertEqual(dec_locked, [4614]) self.assertIsNone(node.cp_hicache) self.assertEqual(node.host_len, 0) def test_remove_undrained_write_acks_splits_shared_ack(self): """A grouped ack covering several nodes only loses the rolled-back node id; other nodes' completion is preserved.""" cache = HiRadixCache.__new__(HiRadixCache) cache.cache_controller = types.SimpleNamespace( ack_write_queue=[ HiCacheAck(MagicMock(), MagicMock(), [4614, 4615]), ], ) removed = cache._remove_undrained_write_acks(4614) self.assertTrue(removed) self.assertEqual( [ack.node_ids for ack in cache.cache_controller.ack_write_queue], [[4615]], ) self.assertFalse(cache._remove_undrained_write_acks(9999)) def test_writing_check_ignores_duplicate_ready_ack_for_same_node(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 1 cache.enable_storage = False cache.pending_host_backups = {} class ReadyEvent: def query(self): return True def synchronize(self): pass node = TreeNode() node.id = 4614 dec_locked = [] cache.ongoing_write_through = {4614: node} cache.dec_node_lock_ref = lambda released_node: dec_locked.append(released_node.id) cache.cache_controller = types.SimpleNamespace( ack_write_queue=[ HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]), HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]), ], ) cache.writing_check() self.assertEqual(cache.ongoing_write_through, {}) self.assertEqual(cache.cache_controller.ack_write_queue, []) self.assertEqual(dec_locked, [4614]) def test_writing_check_drops_stale_duplicate_ack_before_valid_ack(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 1 cache.enable_storage = False cache.pending_host_backups = {} cache._completed_write_ack_ids = {4614} class ReadyEvent: def query(self): return True def synchronize(self): pass node = TreeNode() node.id = 4615 dec_locked = [] cache.ongoing_write_through = {4615: node} cache.dec_node_lock_ref = lambda released_node: dec_locked.append(released_node.id) cache.cache_controller = types.SimpleNamespace( ack_write_queue=[ HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]), HiCacheAck(ReadyEvent(), ReadyEvent(), [4615]), ], ) cache.writing_check() self.assertEqual(cache.ongoing_write_through, {}) self.assertEqual(cache.cache_controller.ack_write_queue, []) self.assertEqual(dec_locked, [4615]) self.assertEqual(cache._completed_write_ack_ids, set()) def test_pending_split_probe_does_not_run_collective_writing_check(self): """Pending-split probing runs inside radix insert, where only a subset of CP ranks may take this branch. It must not opportunistically drain write acks via writing_check(), because writing_check can issue a TP collective and race against the scheduler's inflight-poll collective on other ranks. """ cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 2 cache.tp_group = object() cache.enable_storage = False cache.pending_host_backups = {} node = TreeNode() node.id = 10616 cache.ongoing_write_through = {10616: node} cache.dec_node_lock_ref = lambda released_node: None class ReadyEvent: def query(self): return True def synchronize(self): pass cache.cache_controller = types.SimpleNamespace( ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [10616])], ) cache._cp_hicache_all_reduce = lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("pending split probe must not issue a collective") ) cache._cp_drain_write_acks_before_pending_split() self.assertEqual(cache.ongoing_write_through, {10616: node}) self.assertEqual( [ack.node_ids for ack in cache.cache_controller.ack_write_queue], [[10616]], ) def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = FakeReserveWriteController( [ HiCacheWriteFailure(required_host_slots=2), HiCacheWriteFailure(required_host_slots=2), ] ) cache.evictable_host_leaves = set() cache.eviction_strategy = FakeEvictionStrategy() cache._record_remove_event = lambda node: None cache.ongoing_write_through = {} cache.pending_host_backups = {} node = TreeNode() node.id = 124 node.value = torch.arange(16, dtype=torch.int64) with self.assertRaisesRegex( RuntimeError, "owner-lane admission predicted no host deficit", ): cache.write_backup(node) self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) self.assertEqual(cache.pending_host_backups, {}) self.assertEqual(cache.cache_controller.submitted, []) 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 = FakeReserveWriteController( [lambda device_indices: HiCacheWriteReservation( metadata=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, ), host_indices=torch.empty((0,), dtype=torch.int64), physical_device_indices=torch.empty((0,), dtype=torch.int64), node_id=123, )] ) cache.ongoing_write_through = {} cache.pending_host_backups = {} 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, 0) self.assertIsNone(node.cp_hicache) self.assertIn(node.id, cache.pending_host_backups) self.assertEqual(len(cache.cache_controller.submitted), 1) 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]) def test_cp_owner_token_counts_use_page_owners_and_page_size(self): cache = HiRadixCache.__new__(HiRadixCache) cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=4, cp_rank=0) ) counts = cache._cp_owner_token_counts( torch.tensor([0, 1, 1, 2], dtype=torch.int8), page_size=64, ) self.assertEqual(counts, (64, 128, 64, 0)) def test_cp_metadata_count_asserts_local_owner_lane(self): cache = HiRadixCache.__new__(HiRadixCache) cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=3, cp_rank=1), has_draft_hicache=True, ) metadata = CpHiCacheNodeMetadata( logical_len=192, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(100, 164, dtype=torch.int64), draft_host_indices=torch.arange(200, 264, dtype=torch.int64), page_owners=torch.tensor([0, 1, 2], dtype=torch.int8), page_size=64, ) counts = cache._cp_assert_metadata_counts(metadata, context="unit") self.assertEqual(counts, (64, 64, 64)) def test_cp_metadata_count_asserts_mismatched_local_slots(self): cache = HiRadixCache.__new__(HiRadixCache) cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=1), has_draft_hicache=False, ) metadata = CpHiCacheNodeMetadata( logical_len=128, owned_positions=torch.empty((0,), dtype=torch.int64), host_indices=torch.empty((0,), dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=64, ) with self.assertRaisesRegex(RuntimeError, "owner lane mismatch"): cache._cp_assert_metadata_counts(metadata, context="unit") def test_cp_host_capacity_snapshot_counts_committed_and_pending_draft(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0), has_draft_hicache=True, draft_mem_pool_host=types.SimpleNamespace(size=512), ) cache.token_to_kv_pool_host = types.SimpleNamespace(size=512) cache.pending_host_backups = {} root = TreeNode() root.children = {} cache.root_node = root committed = TreeNode() committed.id = 501 committed.parent = root committed.key = RadixKey([1]) committed.host_len = 128 committed.cp_hicache = CpHiCacheNodeMetadata( logical_len=128, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(64, dtype=torch.int64), draft_host_indices=torch.arange(100, 164, dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=64, ) root.children[1] = committed pending = TreeNode() pending.id = 502 pending_metadata = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(200, 264, dtype=torch.int64), draft_host_indices=torch.arange(300, 364, dtype=torch.int64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) cache.pending_host_backups[pending.id] = PendingHiCacheBackup( node=pending, metadata=pending_metadata, logical_len=64, ) snapshot = cache._cp_host_capacity_snapshot() self.assertEqual(snapshot.target_capacity, (512, 512)) self.assertEqual(snapshot.draft_capacity, (512, 512)) self.assertEqual(snapshot.committed_target, (64, 64)) self.assertEqual(snapshot.committed_draft, (64, 64)) self.assertEqual(snapshot.pending_target, (64, 0)) self.assertEqual(snapshot.pending_draft, (64, 0)) def test_cp_write_admission_uses_padded_owner_vector_for_valid_tail(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.page_size = 64 cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=8, cp_rank=0, page_size=64), has_draft_hicache=True, draft_mem_pool_host=types.SimpleNamespace(size=1024), ) cache.token_to_kv_pool_host = types.SimpleNamespace(size=1024) cache.pending_host_backups = {} cache.root_node = TreeNode() cache.root_node.children = {} cache.evictable_host_leaves = set() admission = cache._cp_build_write_admission( torch.arange(64, 164, dtype=torch.int64), node_id=503, phase="unit", ) self.assertEqual(admission.required_by_owner, (64, 64, 0, 0, 0, 0, 0, 0)) self.assertEqual( admission.target_available_by_owner, (1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024), ) self.assertEqual( admission.target_available_by_owner, admission.draft_available_by_owner ) self.assertEqual(admission.deficit_by_owner, (0, 0, 0, 0, 0, 0, 0, 0)) def test_cp_capacity_snapshot_counts_padded_valid_tail_for_target_and_draft(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0, page_size=64), has_draft_hicache=True, draft_mem_pool_host=types.SimpleNamespace(size=1024), ) cache.token_to_kv_pool_host = types.SimpleNamespace(size=1024) cache.pending_host_backups = {} root = TreeNode() root.children = {} cache.root_node = root committed = TreeNode() committed.id = 504 committed.parent = root committed.key = RadixKey(list(range(100))) committed.host_len = 100 committed.cp_hicache = CpHiCacheNodeMetadata( logical_len=100, padded_len=128, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(100, 164, dtype=torch.int64), draft_host_indices=torch.arange(200, 264, dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=64, ) root.children[0] = committed snapshot = cache._cp_host_capacity_snapshot() self.assertEqual(snapshot.committed_target, (64, 64)) self.assertEqual(snapshot.committed_draft, (64, 64)) 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_leaf_status_skips_device_valid_host_backed_node(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.evictable_host_leaves = set() node = TreeNode() 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, ) cache.evictable_host_leaves.add(node) cache._update_host_leaf_status(node) self.assertNotIn(node, cache.evictable_host_leaves) 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) def test_cp_host_eviction_plan_is_stable_across_leaf_insertion_order(self): def make_cache(order): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0), has_draft_hicache=False, ) cache.eviction_strategy = FakeEvictionStrategy() cache.pending_host_backups = {} cache.root_node = TreeNode() cache.root_node.children = {} cache.get_child_key_fn = lambda key: key.token_ids[0] cache._clear_pin = lambda node: None nodes = [] for node_id in (42, 41): node = TreeNode(id=node_id) node.parent = cache.root_node node.key = RadixKey([node_id]) node.value = None node.lock_ref = 0 node.host_ref_counter = 0 node.host_len = 64 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(node_id * 100, node_id * 100 + 64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) cache.root_node.children[node_id] = node nodes.append(node) cache.evictable_host_leaves = {nodes[i] for i in order} return cache first = make_cache([0, 1])._plan_cp_host_evictions((64, 0)) second = make_cache([1, 0])._plan_cp_host_evictions((64, 0)) self.assertEqual([node.id for node in first.victims], [41]) self.assertEqual([node.id for node in second.victims], [41]) self.assertEqual(first.planned_freed, (64, 0)) self.assertEqual(first.remaining_deficit, (0, 0)) def test_cp_host_eviction_plan_skips_pending_backup_nodes(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = types.SimpleNamespace( cp_shared_kv_layout=FakeCpLayout(cp_size=1, cp_rank=0), has_draft_hicache=False, ) cache.eviction_strategy = FakeEvictionStrategy() cache.root_node = TreeNode() cache.root_node.children = {} cache.get_child_key_fn = lambda key: key.token_ids[0] cache.pending_host_backups = {} cache.evictable_host_leaves = set() for node_id in (10, 11): node = TreeNode(id=node_id) node.parent = cache.root_node node.key = RadixKey([node_id]) node.value = None node.lock_ref = 0 node.host_ref_counter = 0 node.host_len = 64 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=64, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(node_id * 100, node_id * 100 + 64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) cache.root_node.children[node_id] = node cache.evictable_host_leaves.add(node) if node_id == 10: cache.pending_host_backups[node.id] = PendingHiCacheBackup( node=node, metadata=node.cp_hicache, logical_len=64, ) plan = cache._plan_cp_host_evictions((64,)) self.assertEqual([node.id for node in plan.victims], [11]) self.assertEqual(plan.planned_freed, (64,)) class TestHiRadixCacheCPLoadBack(CustomTestCase): def test_cp_load_back_plan_reports_l1_free_room_without_blocking_exact_fit(self): class FreeRoomAllocator: def compute_owner_lane_stats(self, _page_owners): return [1, 0], [0, 8], [1, 0] def compute_owner_lane_capacity_pages(self): return [8, 8] cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.page_size = 64 cache.hicache_l1_free_room_ratio = 0.5 cache.hicache_l1_free_room_trigger_ratio = 0.25 cache.token_to_kv_pool_allocator = FreeRoomAllocator() node = TreeNode(id=701) node.host_len = 64 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=64, padded_len=64, owned_positions=torch.arange(64, dtype=torch.int64), host_indices=torch.arange(64, dtype=torch.int64), page_owners=torch.tensor([0], dtype=torch.int8), page_size=64, ) plan = cache._build_cp_load_back_plan([node], node_id=701) self.assertEqual(plan.required_by_owner, [1, 0]) self.assertEqual(plan.available_by_owner, [0, 8]) # Synchronous load-back admission must only block on exact capacity. self.assertEqual(plan.deficit_by_owner, [1, 0]) # The free-room target is still reported for observability/proactive policy. # required=1 page, available=0, target_room=ceil(8*0.5)=4 pages. self.assertEqual(plan.free_room_deficit_by_owner, [4, 0]) 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.page_size = 1 cache.token_to_kv_pool_allocator = FakeTokenAllocator() 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.page_size = 1 cache.token_to_kv_pool_allocator = FakeTokenAllocator() 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_device_valid_host_missing_remains_device_hit(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 = torch.arange(8, dtype=torch.int64) node.host_value = None node.host_len = 0 node.cp_hicache = None root.children[0] = 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_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_defers_split_when_child_backup_pending(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.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) 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 child = TreeNode() child.id = 130 child.parent = root child.key = RadixKey(list(range(8))) child.value = torch.arange(8, dtype=torch.int64) child.host_len = 0 metadata = CpHiCacheNodeMetadata( logical_len=8, owned_positions=torch.arange(8, dtype=torch.int64), host_indices=torch.arange(50, 58, dtype=torch.int64), page_owners=torch.zeros(8, dtype=torch.int8), page_size=1, ) cache.pending_host_backups[child.id] = PendingHiCacheBackup( node=child, metadata=metadata, logical_len=8 ) root.children[0] = child result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 9]))) self.assertIs(result.pending_backup_deferred_node, child) self.assertIs(root.children[0], child) self.assertEqual(child.key.token_ids, list(range(8))) self.assertEqual(result.device_indices.tolist(), []) self.assertIs(result.last_device_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_cp_match_prefix_reports_valid_tail_host_hit(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 4 cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) 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 = 140 node.parent = root node.key = RadixKey(list(range(6))) node.value = None node.host_value = None node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, padded_len=8, owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) root.children[(0, 1, 2, 3)] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6))))) self.assertEqual(result.device_indices.tolist(), []) # CP HiCache exposes page-granular host hits. A non-page exact tail is # floored and the tail is recomputed by the incoming request. self.assertEqual(result.host_hit_length, 4) self.assertIs(result.last_device_node, root) self.assertEqual(result.last_host_node.key.token_ids, [0, 1, 2, 3]) self.assertIsNot(result.last_host_node, node) def test_cp_backed_tail_split_floors_to_page_boundary(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 4 cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) 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 root.children = {} cache.root_node = root node = TreeNode() node.id = 141 node.parent = root node.key = RadixKey(list(range(6))) node.value = None node.host_value = None node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, padded_len=8, owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) root.children[(0, 1, 2, 3)] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 3, 4]))) self.assertEqual(result.device_indices.tolist(), []) self.assertEqual(result.host_hit_length, 4) self.assertIs(result.last_device_node, root) self.assertEqual(result.last_host_node.key.token_ids, [0, 1, 2, 3]) self.assertEqual(result.last_host_node.host_len, 4) self.assertEqual(result.last_host_node.cp_hicache.logical_len, 4) self.assertEqual(result.last_host_node.cp_hicache.padded_len, 4) self.assertNotIn((4, 5), result.last_host_node.children) def test_cp_backed_tail_split_before_first_page_returns_root_match(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 4 cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) 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 root.children = {} cache.root_node = root node = TreeNode() node.id = 142 node.parent = root node.key = RadixKey(list(range(6))) node.value = None node.host_value = None node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, padded_len=8, owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) root.children[(0, 1, 2, 3)] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2]))) self.assertEqual(result.device_indices.tolist(), []) self.assertEqual(result.host_hit_length, 0) self.assertIs(result.last_device_node, root) self.assertIs(result.last_host_node, root) self.assertIs(root.children[(0, 1, 2, 3)], node) self.assertEqual(node.key.token_ids, list(range(6))) def test_cp_prepare_probe_floors_backed_tail_partial_hit_to_page(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.disable = False cache.page_size = 4 cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) root = TreeNode() root.key = RadixKey([]) root.children = {} cache.root_node = root node = TreeNode() node.id = 143 node.parent = root node.key = RadixKey(list(range(6))) node.value = None node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, padded_len=8, owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) root.children[(0, 1, 2, 3)] = node prefix_len = cache._probe_existing_radix_prefix_len_no_split( RadixKey([0, 1, 2, 3, 4]) ) self.assertEqual(prefix_len, 4) def test_cp_prepare_probe_floors_exact_valid_tail_when_request_extends(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.disable = False cache.page_size = 4 cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) root = TreeNode() root.key = RadixKey([]) root.children = {} cache.root_node = root node = TreeNode() node.id = 145 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, padded_len=8, owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) root.children[(0, 1, 2, 3)] = node prefix_len = cache._probe_existing_radix_prefix_len_no_split( RadixKey(list(range(10))) ) self.assertEqual(prefix_len, 4) def test_cp_match_prefix_floors_exact_valid_tail_when_request_extends(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 4 cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key: (key, None) cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache.enable_storage = False cache.enable_kv_cache_events = False root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.host_len = 0 root.children = {} cache.root_node = root node = TreeNode() node.id = 146 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_value = None node.host_len = 0 node.cp_hicache = None root.children[(0, 1, 2, 3)] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(10))))) self.assertEqual(result.device_indices.tolist(), [0, 1, 2, 3]) self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3]) self.assertNotIn((4, 5), result.last_device_node.children) def test_cp_match_prefix_floors_exact_valid_tail_for_exact_key(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.device = "cpu" cache.disable = False cache.page_size = 4 cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key: (key, None) cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache.enable_storage = False cache.enable_kv_cache_events = False root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.host_len = 0 root.children = {} cache.root_node = root node = TreeNode() node.id = 149 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_value = None node.host_len = 0 node.cp_hicache = None root.children[(0, 1, 2, 3)] = node result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6))))) self.assertEqual(result.device_indices.tolist(), [0, 1, 2, 3]) self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3]) self.assertNotIn((4, 5), result.last_device_node.children) def test_cp_insert_floors_backed_tail_split_to_page_boundary(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.disable = False cache.is_eagle = False cache.page_size = 4 cache.pending_host_backups = {} cache.cache_controller = types.SimpleNamespace( has_draft_hicache=False, write_policy="write_back", ) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache._inc_hit_count = lambda *args, **kwargs: None cache._record_store_event = lambda node: None cache.evictable_size_ = 0 cache.enable_storage = False cache.enable_kv_cache_events = False root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.children = {} cache.root_node = root node = TreeNode() node.id = 144 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_len = 6 node.cp_hicache = CpHiCacheNodeMetadata( logical_len=6, padded_len=8, owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), page_owners=torch.tensor([0, 1], dtype=torch.int8), page_size=4, ) root.children[(0, 1, 2, 3)] = node result = cache.insert( InsertParams( key=RadixKey([0, 1, 2, 3, 4]), value=torch.arange(5, dtype=torch.int64), ) ) self.assertEqual(result.prefix_len, 4) parent = root.children[(0, 1, 2, 3)] self.assertEqual(parent.key.token_ids, [0, 1, 2, 3]) self.assertEqual(parent.cp_hicache.logical_len, 4) self.assertEqual(parent.cp_hicache.padded_len, 4) self.assertNotIn((4, 5), parent.children) new_tail = parent.children[(4,)] self.assertEqual(new_tail.key.token_ids, [4]) self.assertEqual(new_tail.value.tolist(), [4]) def test_cp_insert_defers_stale_tail_prune_while_backup_pending(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.disable = False cache.is_eagle = False cache.page_size = 4 cache.pending_host_backups = {} cache.ongoing_write_through = {} cache.cache_controller = types.SimpleNamespace( has_draft_hicache=False, write_policy="write_through", ack_write_queue=[], ) cache.tp_world_size = 1 cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache._inc_hit_count = lambda *args, **kwargs: None cache._record_store_event = lambda node: None cache.evictable_size_ = 0 cache.enable_storage = False cache.enable_kv_cache_events = False root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.children = {} cache.root_node = root node = TreeNode() node.id = 152 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_len = 0 node.cp_hicache = None root.children[(0, 1, 2, 3)] = node cache.ongoing_write_through[node.id] = node result = cache.insert( InsertParams( key=RadixKey(list(range(10))), value=torch.arange(10, dtype=torch.int64), ) ) self.assertEqual(result.prefix_len, 0) self.assertIs(result.pending_backup_deferred_node, node) self.assertIs(root.children[(0, 1, 2, 3)], node) self.assertEqual(node.key.token_ids, list(range(6))) self.assertEqual(node.value.tolist(), list(range(6))) def test_cp_insert_extends_from_page_boundary_after_exact_valid_tail(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.disable = False cache.is_eagle = False cache.page_size = 4 cache.pending_host_backups = {} cache.ongoing_write_through = {} cache.cache_controller = types.SimpleNamespace( has_draft_hicache=False, write_policy="write_through", ) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache._inc_hit_count = lambda *args, **kwargs: None cache._record_store_event = lambda node: None cache.evictable_size_ = 0 cache.protected_size_ = 0 cache.enable_storage = False cache.enable_kv_cache_events = False cache.inc_node_lock_ref = lambda node: None root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.children = {} cache.root_node = root node = TreeNode() node.id = 147 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_len = 0 node.cp_hicache = None root.children[(0, 1, 2, 3)] = node reservation = make_write_reservation( torch.arange(4, 10, dtype=torch.int64), node_id=148, host_start=190 ) prepared = PreparedCpHiCacheBackup( node_id=148, reservation=reservation, metadata=reservation.metadata, logical_len=6, ) result = cache.insert( InsertParams( key=RadixKey(list(range(10))), value=torch.arange(10, dtype=torch.int64), cp_hicache_prepared_backup=prepared, ) ) self.assertEqual(result.prefix_len, 4) parent = root.children[(0, 1, 2, 3)] self.assertEqual(parent.key.token_ids, [0, 1, 2, 3]) self.assertNotIn((4, 5), parent.children) new_tail = parent.children[(4, 5, 6, 7)] self.assertEqual(new_tail.key.token_ids, [4, 5, 6, 7, 8, 9]) self.assertTrue(prepared.attached) self.assertIs(cache.pending_host_backups[148].node, new_tail) def test_cp_insert_replaces_exact_valid_tail_from_page_boundary(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.disable = False cache.is_eagle = False cache.page_size = 4 cache.pending_host_backups = {} cache.ongoing_write_through = {} cache.cache_controller = types.SimpleNamespace( has_draft_hicache=False, write_policy="write_through", ) cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) cache.key_match_fn = lambda child_key, key: _key_match_paged( child_key, key, page_size=4 ) cache.maybe_bigram_convert = lambda key, value=None: (key, value) cache._update_leaf_status = lambda node: None cache._update_host_leaf_status = lambda node: None cache._inc_hit_count = lambda *args, **kwargs: None cache._record_store_event = lambda node: None cache.evictable_size_ = 0 cache.protected_size_ = 0 cache.enable_storage = False cache.enable_kv_cache_events = False cache.inc_node_lock_ref = lambda node: None root = TreeNode() root.key = RadixKey([]) root.value = torch.empty((0,), dtype=torch.int64) root.children = {} cache.root_node = root node = TreeNode() node.id = 150 node.parent = root node.key = RadixKey(list(range(6))) node.value = torch.arange(6, dtype=torch.int64) node.host_len = 0 node.cp_hicache = None root.children[(0, 1, 2, 3)] = node reservation = make_write_reservation( torch.arange(4, 6, dtype=torch.int64), node_id=151, host_start=210 ) prepared = PreparedCpHiCacheBackup( node_id=151, reservation=reservation, metadata=reservation.metadata, logical_len=2, ) result = cache.insert( InsertParams( key=RadixKey(list(range(6))), value=torch.arange(6, dtype=torch.int64), cp_hicache_prepared_backup=prepared, ) ) self.assertEqual(result.prefix_len, 4) parent = root.children[(0, 1, 2, 3)] self.assertEqual(parent.key.token_ids, [0, 1, 2, 3]) new_tail = parent.children[(4, 5)] self.assertIsNot(new_tail, node) self.assertEqual(new_tail.key.token_ids, [4, 5]) self.assertTrue(prepared.attached) self.assertIs(cache.pending_host_backups[151].node, new_tail) 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, {}) class TestCPHiCacheLayerBackupNotifications(CustomTestCase): def _call_store_index_fast_path(self, out_loc): from sglang.srt.layers.attention.nsa import nsa_indexer notifications = [] fused_calls = [] class FakePool: page_size = 64 start_layer = 0 def get_index_k_with_scale_buffer(self, layer_id): return torch.empty((1,), dtype=torch.uint8) def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"): notifications.append((layer_id, source)) forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool()) indexer = types.SimpleNamespace() def fake_fused_store(key, buf, loc, page_size): fused_calls.append((key, buf, loc.clone(), page_size)) with ( patch.object(nsa_indexer, "_is_cuda", True), patch.object(nsa_indexer, "_is_fp8_fnuz", False), patch.object(nsa_indexer, "can_use_nsa_fused_store", return_value=True), patch.object( nsa_indexer, "fused_store_index_k_cache", side_effect=fake_fused_store ), ): nsa_indexer.Indexer._store_index_k_cache( indexer, forward_batch, layer_id=5, key=torch.empty((max(out_loc.numel(), 1), 128), dtype=torch.float32), out_loc_override=out_loc, ) return notifications, fused_calls def test_nsa_indexer_fused_store_does_not_notify_cp_hicache_layer_backup(self): notifications, fused_calls = self._call_store_index_fast_path( torch.tensor([1, 2, 3], dtype=torch.int64) ) self.assertEqual(notifications, []) self.assertEqual(len(fused_calls), 1) self.assertEqual(fused_calls[0][2].tolist(), [1, 2, 3]) def test_nsa_indexer_empty_store_does_not_notify_cp_hicache_layer_backup(self): notifications, fused_calls = self._call_store_index_fast_path( torch.empty((0,), dtype=torch.int64) ) self.assertEqual(notifications, []) self.assertEqual(fused_calls, []) def test_cp_shared_zero_local_index_store_does_not_notify_layer_backup(self): from sglang.srt.layers.attention.nsa import nsa_indexer notifications = [] class FakePool: page_size = 64 def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"): notifications.append((layer_id, source)) forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool()) indexer = types.SimpleNamespace(nsa_enable_prefill_cp=True) with ( patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True), patch.object( nsa_indexer, "get_cp_shared_kv_local_out_cache_loc", return_value=torch.empty((0,), dtype=torch.int64), ), patch.object( nsa_indexer, "get_cp_shared_kv_local_physical_out_cache_loc", side_effect=AssertionError("must not require physical locs"), ), ): handled = nsa_indexer.Indexer._store_cp_shared_local_index_k_cache( indexer, forward_batch=forward_batch, layer_id=6, local_key=torch.empty((0, 128), dtype=torch.float32), act_quant=None, ) self.assertTrue(handled) self.assertEqual(notifications, []) if __name__ == "__main__": unittest.main()