sglang/test/registered/unit/mem_cache/test_cp_hicache_metadata.py

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