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

import sys
import types
import unittest
from unittest.mock import patch

import torch

# Prefer the real sgl_kernel package when the test image provides it so custom
# Torch operators are registered. Fall back to stubs on local CPU-only hosts.
try:
    import sgl_kernel  # noqa: F401
    import sgl_kernel.kvcacheio  # noqa: F401
except (ImportError, RuntimeError):
    if "sgl_kernel" not in sys.modules:
        sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
    sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
    sys.modules["sgl_kernel"].__path__ = []
    if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):

        def _sgl_kernel_getattr(name):
            if name.startswith("__"):
                raise AttributeError(name)
            fn = lambda *args, **kwargs: None
            setattr(sys.modules["sgl_kernel"], name, fn)
            return fn

        sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
    if "sgl_kernel.quantization" not in sys.modules:
        quantization_stub = types.ModuleType("sgl_kernel.quantization")
        quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"

        def _quantization_getattr(name):
            if name.startswith("__"):
                raise AttributeError(name)
            fn = lambda *args, **kwargs: None
            setattr(quantization_stub, name, fn)
            return fn

        quantization_stub.__getattr__ = _quantization_getattr
        for _name in (
            "ggml_dequantize",
            "ggml_moe_a8",
            "ggml_moe_a8_vec",
            "ggml_moe_get_block_size",
            "ggml_mul_mat_a8",
            "ggml_mul_mat_vec_a8",
        ):
            setattr(quantization_stub, _name, lambda *args, **kwargs: None)
        sys.modules["sgl_kernel.quantization"] = quantization_stub
    for _name in (
        "sgl_per_token_group_quant_8bit",
        "sgl_per_token_group_quant_fp8",
        "sgl_per_token_quant_fp8",
        "fp8_blockwise_scaled_mm",
        "fp8_scaled_mm",
        "silu_and_mul",
    ):
        if not hasattr(sys.modules["sgl_kernel"], _name):
            setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None)
    _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
    for _schema in (
        "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
        "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
        "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
        "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
        "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
    ):
        try:
            _sgl_kernel_lib.define(_schema)
        except RuntimeError as exc:
            if (
                "already" not in str(exc).lower()
                and "duplicate" not in str(exc).lower()
            ):
                raise
if "sgl_kernel.kvcacheio" not in sys.modules:
    sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")

_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
for _schema in (
    "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
    "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
    "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
    "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
    "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
):
    try:
        _sgl_kernel_lib.define(_schema)
    except RuntimeError as exc:
        if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
            raise

from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=2, suite="stage-a-test-cpu")


class TestCpHiCacheImports(CustomTestCase):
    def test_cp_hicache_public_imports_without_sgl_kernel(self):
        import subprocess

        subprocess.run(
            [
                sys.executable,
                "-c",
                "from sglang.srt.mem_cache.hiradix_cache import HiRadixCache, CpHiCacheNodeMetadata; "
                "from sglang.srt.managers.cache_controller import HiCacheController; "
                "print('OK')",
            ],
            check=True,
            capture_output=True,
            text=True,
        )


class TestCpHiCacheNodeMetadata(CustomTestCase):
    def test_split_zero_len_moves_all_positions_to_child(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([1, 3, 7], dtype=torch.int64),
            host_indices=torch.tensor([10, 11, 12], dtype=torch.int64),
            page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
            page_size=1,
        )

        parent, child = metadata.split(0)

        self.assertEqual(parent.logical_len, 0)
        self.assertEqual(parent.owned_positions.tolist(), [])
        self.assertEqual(parent.host_indices.tolist(), [])
        self.assertEqual(child.logical_len, 8)
        self.assertEqual(child.owned_positions.tolist(), [1, 3, 7])
        self.assertEqual(child.host_indices.tolist(), [10, 11, 12])

    def test_split_rebases_child_positions(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=10,
            owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
            host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
            page_owners=torch.zeros(max(10, 0), dtype=torch.int8),
            page_size=1,
        )

        parent, child = metadata.split(5)

        self.assertEqual(parent.logical_len, 5)
        self.assertEqual(parent.owned_positions.tolist(), [0, 2])
        self.assertEqual(parent.host_indices.tolist(), [20, 21])
        self.assertEqual(child.logical_len, 5)
        self.assertEqual(child.owned_positions.tolist(), [0, 4])
        self.assertEqual(child.host_indices.tolist(), [22, 23])

    def test_split_keeps_draft_host_indices_aligned_with_owned_positions(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=10,
            owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
            host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
            draft_host_indices=torch.tensor([120, 121, 122, 123], dtype=torch.int64),
            page_owners=torch.zeros(max(10, 0), dtype=torch.int8),
            page_size=1,
        )

        parent, child = metadata.split(5)

        self.assertEqual(parent.host_indices.tolist(), [20, 21])
        self.assertEqual(parent.draft_host_indices.tolist(), [120, 121])
        self.assertEqual(child.owned_positions.tolist(), [0, 4])
        self.assertEqual(child.host_indices.tolist(), [22, 23])
        self.assertEqual(child.draft_host_indices.tolist(), [122, 123])

    def test_zero_owned_metadata_is_valid(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=64,
            owned_positions=torch.empty((0,), dtype=torch.int32),
            host_indices=torch.empty((0,), dtype=torch.int32),
            page_owners=torch.zeros(max(64, 0), dtype=torch.int8),
            page_size=1,
        )

        self.assertEqual(metadata.logical_len, 64)
        self.assertEqual(metadata.owned_positions.device.type, "cpu")
        self.assertEqual(metadata.host_indices.device.type, "cpu")
        self.assertEqual(metadata.owned_positions.dtype, torch.int64)
        self.assertEqual(metadata.host_indices.dtype, torch.int64)

    def test_non_int64_inputs_are_converted(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([1, 3], dtype=torch.int32),
            host_indices=torch.tensor([10, 11], dtype=torch.int32),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )

        self.assertEqual(metadata.owned_positions.dtype, torch.int64)
        self.assertEqual(metadata.host_indices.dtype, torch.int64)

    def test_negative_logical_len_raises(self):
        with self.assertRaisesRegex(ValueError, "logical_len"):
            CpHiCacheNodeMetadata(
                logical_len=-1,
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
                page_owners=torch.zeros(0, dtype=torch.int8),
                page_size=1,
            )

    def test_metadata_does_not_alias_input_tensors(self):
        owned_positions = torch.tensor([1, 3], dtype=torch.int64)
        host_indices = torch.tensor([10, 11], dtype=torch.int64)

        metadata = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=owned_positions,
            host_indices=host_indices,
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        owned_positions[0] = 2
        host_indices[0] = 12

        self.assertEqual(metadata.owned_positions.tolist(), [1, 3])
        self.assertEqual(metadata.host_indices.tolist(), [10, 11])

    def test_invalid_split_raises(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([1], dtype=torch.int64),
            host_indices=torch.tensor([9], dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )

        with self.assertRaisesRegex(ValueError, "split_len"):
            metadata.split(5)

    def test_unsorted_positions_raise(self):
        with self.assertRaisesRegex(ValueError, "sorted"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([2, 1], dtype=torch.int64),
                host_indices=torch.tensor([9, 10], dtype=torch.int64),
                page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
                page_size=1,
            )

    def test_duplicate_positions_raise(self):
        with self.assertRaisesRegex(ValueError, "strictly increasing"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([1, 1], dtype=torch.int64),
                host_indices=torch.tensor([9, 10], dtype=torch.int64),
                page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
                page_size=1,
            )

    def test_length_mismatch_raises(self):
        with self.assertRaisesRegex(ValueError, "same length"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([1, 2], dtype=torch.int64),
                host_indices=torch.tensor([9], dtype=torch.int64),
                page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
                page_size=1,
            )

    def test_draft_host_length_mismatch_raises(self):
        with self.assertRaisesRegex(ValueError, "draft_host_indices.*same length"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([1, 2], dtype=torch.int64),
                host_indices=torch.tensor([9, 10], dtype=torch.int64),
                draft_host_indices=torch.tensor([109], dtype=torch.int64),
                page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
                page_size=1,
            )

    def test_out_of_range_positions_raise(self):
        with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([4], dtype=torch.int64),
                host_indices=torch.tensor([9], dtype=torch.int64),
                page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
                page_size=1,
            )

    # ── New: validators for page_owners + page_size (the fields that carry the
    # CP owner pattern across a HiCache write→load round-trip).

    def test_zero_page_size_raises(self):
        with self.assertRaisesRegex(ValueError, "page_size must be positive"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
                page_owners=torch.empty((0,), dtype=torch.int8),
                page_size=0,
            )

    def test_logical_len_not_multiple_of_page_size_raises(self):
        with self.assertRaisesRegex(ValueError, "multiple of"):
            CpHiCacheNodeMetadata(
                logical_len=10,
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
                page_owners=torch.zeros(2, dtype=torch.int8),
                page_size=4,  # 10 % 4 != 0
            )

    def test_page_owners_length_mismatch_raises(self):
        with self.assertRaisesRegex(ValueError, "page_owners length"):
            CpHiCacheNodeMetadata(
                logical_len=8,
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
                page_owners=torch.zeros(3, dtype=torch.int8),  # expected 8/4=2
                page_size=4,
            )

    def test_negative_page_owners_raises(self):
        with self.assertRaisesRegex(ValueError, "page_owners.*non-negative"):
            CpHiCacheNodeMetadata(
                logical_len=8,
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
                page_owners=torch.tensor([-1, 0], dtype=torch.int8),
                page_size=4,
            )

    def test_page_owners_normalized_to_int8_cpu(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.empty((0,), dtype=torch.int64),
            host_indices=torch.empty((0,), dtype=torch.int64),
            # Pass int64 to test normalization.
            page_owners=torch.tensor([0, 1], dtype=torch.int64),
            page_size=4,
        )
        self.assertEqual(metadata.page_owners.dtype, torch.int8)
        self.assertEqual(metadata.page_owners.device.type, "cpu")
        self.assertEqual(metadata.page_owners.tolist(), [0, 1])

    def test_page_owners_not_aliased_to_input(self):
        page_owners = torch.tensor([0, 1, 0, 1], dtype=torch.int8)
        metadata = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.empty((0,), dtype=torch.int64),
            host_indices=torch.empty((0,), dtype=torch.int64),
            page_owners=page_owners,
            page_size=4,
        )
        page_owners[0] = 1
        self.assertEqual(metadata.page_owners.tolist(), [0, 1, 0, 1])

    def test_split_non_page_aligned_split_len_raises(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([0, 8], dtype=torch.int64),
            host_indices=torch.tensor([20, 21], dtype=torch.int64),
            page_owners=torch.zeros(4, dtype=torch.int8),
            page_size=4,
        )
        with self.assertRaisesRegex(ValueError, "must be a multiple of page_size"):
            metadata.split(5)  # 5 % 4 != 0

    def test_split_preserves_page_owners_slice(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=32,  # 8 pages of page_size=4
            owned_positions=torch.tensor([0, 4, 16, 28], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor(
                [0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.int8
            ),
            page_size=4,
        )
        parent, child = metadata.split(16)  # split at page 4 of 8.
        self.assertEqual(parent.page_owners.tolist(), [0, 1, 0, 1])
        self.assertEqual(child.page_owners.tolist(), [0, 1, 0, 1])
        self.assertEqual(parent.page_size, 4)
        self.assertEqual(child.page_size, 4)

    def test_split_at_zero_yields_empty_parent_page_owners(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([0], dtype=torch.int64),
            host_indices=torch.tensor([50], dtype=torch.int64),
            page_owners=torch.tensor([0, 1], dtype=torch.int8),
            page_size=4,
        )
        parent, child = metadata.split(0)
        self.assertEqual(parent.page_owners.numel(), 0)
        self.assertEqual(child.page_owners.tolist(), [0, 1])

    def test_split_at_logical_len_yields_empty_child_page_owners(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([0], dtype=torch.int64),
            host_indices=torch.tensor([50], dtype=torch.int64),
            page_owners=torch.tensor([0, 1], dtype=torch.int8),
            page_size=4,
        )
        parent, child = metadata.split(8)
        self.assertEqual(parent.page_owners.tolist(), [0, 1])
        self.assertEqual(child.page_owners.numel(), 0)


class FakeWriteFailure:
    metadata = None

    def __init__(self, required_host_slots):
        self.required_host_slots = required_host_slots


class FakeWriteSuccess:
    required_host_slots = 0

    def __init__(self, metadata):
        self.metadata = metadata


class FakeWriteController:
    def __init__(self, required_host_slots):
        self.required_host_slots = required_host_slots
        self.calls = 0
        self.write_policy = "write_through"
        self.evicted_host_indices = []

    def write(self, device_indices, node_id=-1, priority=None):
        self.calls += 1
        if self.calls == 1:
            return FakeWriteFailure(self.required_host_slots)
        return FakeWriteSuccess(
            CpHiCacheNodeMetadata(
                logical_len=len(device_indices),
                owned_positions=torch.tensor([0], dtype=torch.int64),
                host_indices=torch.tensor([99], dtype=torch.int64),
                page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8),
                page_size=1,
            )
        )

    def evict_host(self, host_indices):
        self.evicted_host_indices.append(host_indices.clone())
        return len(host_indices)


class FakeZeroOwnedWriteController:
    write_policy = "write_through"

    def write(self, device_indices, node_id=-1, priority=None):
        return FakeWriteSuccess(
            CpHiCacheNodeMetadata(
                logical_len=len(device_indices),
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
                page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8),
                page_size=1,
            )
        )


class FakeEvictionStrategy:
    def get_priority(self, node):
        return 0


class FakeEvictDeviceController:
    write_policy = "write_through"

    def evict_device(self, device_indices):
        return len(device_indices)


class FakeTokenAllocator:
    def available_size(self):
        return 0


class TestHiRadixCacheCPBackup(CustomTestCase):
    def test_node_backuped_uses_cp_metadata(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        node = TreeNode()
        node.host_len = 8
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([1, 2], dtype=torch.int64),
            host_indices=torch.tensor([10, 11], dtype=torch.int64),
            page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
            page_size=1,
        )

        self.assertTrue(cache._node_backuped(node))

    def test_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_write_backup_retries_by_required_physical_slots(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.cache_controller = FakeWriteController(required_host_slots=1)
        cache.evictable_host_leaves = set()
        cache.eviction_strategy = FakeEvictionStrategy()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache._record_remove_event = lambda node: None
        cache.ongoing_write_through = {}
        cache.inc_node_lock_ref = lambda node: None

        root = TreeNode()
        root.key = RadixKey(token_ids=[], extra_key=None)
        root.value = []
        cache.root_node = root
        evictable_node = TreeNode()
        evictable_node.parent = root
        evictable_node.key = RadixKey(token_ids=[1], extra_key=None)
        evictable_node.value = None
        evictable_node.host_len = 4
        evictable_node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([0], dtype=torch.int64),
            host_indices=torch.tensor([55], dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        root.children[1] = evictable_node
        cache.evictable_host_leaves.add(evictable_node)

        node = TreeNode()
        node.value = torch.arange(16, dtype=torch.int64)

        cache.write_backup(node)

        self.assertEqual(
            cache.cache_controller.evicted_host_indices[0].tolist(), [55]
        )
        self.assertEqual(evictable_node.host_len, 0)
        self.assertIsNone(evictable_node.cp_hicache)
        self.assertNotIn(1, root.children)
        self.assertEqual(node.host_len, 16)

    def test_write_backup_cp_success_returns_logical_length_for_zero_owned_rank(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.cache_controller = FakeZeroOwnedWriteController()
        cache.ongoing_write_through = {}
        cache.inc_node_lock_ref = lambda node: None
        node = TreeNode()
        node.id = 123
        node.value = torch.arange(16, dtype=torch.int64)

        backed_len = cache.write_backup(node)

        self.assertEqual(backed_len, 16)
        self.assertEqual(node.host_len, 16)
        self.assertEqual(node.cp_hicache.host_indices.tolist(), [])

    def test_attach_storage_backend_rejects_cp_hicache_without_controller_call(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "attach_storage_backend": lambda *args, **kwargs: (_ for _ in ()).throw(
                    AssertionError("controller attach must not be called")
                )
            },
        )()

        ok, message = cache.attach_storage_backend("mooncake")

        self.assertFalse(ok)
        self.assertIn("CP shared KV", message)
        self.assertIn("storage backend", message)

    def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.cache_controller = FakeEvictDeviceController()
        cache.evictable_leaves = set()
        cache.evictable_host_leaves = set()
        cache.eviction_strategy = FakeEvictionStrategy()
        cache.token_to_kv_pool_allocator = FakeTokenAllocator()
        cache.evictable_size_ = 4
        cache.protected_size_ = 0
        cache._record_remove_event = lambda node: (_ for _ in ()).throw(
            AssertionError("must not delete backed radix child")
        )

        root = TreeNode()
        root.key = RadixKey(token_ids=[], extra_key=None)
        root.value = []
        cache.root_node = root
        node = TreeNode()
        node.parent = root
        node.key = RadixKey(token_ids=[1, 2, 3, 4], extra_key=None)
        node.value = torch.arange(4, dtype=torch.int64)
        node.host_len = 4
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([0], dtype=torch.int64),
            host_indices=torch.tensor([55], dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        root.children[1] = node
        cache.evictable_leaves.add(node)

        cache.evict(EvictParams(num_tokens=4))

        self.assertIn(1, root.children)
        self.assertIsNone(node.value)
        self.assertIsNotNone(node.cp_hicache)
        self.assertEqual(node.cp_hicache.host_indices.tolist(), [55])


class TestHiRadixCacheCPSplitEvict(CustomTestCase):
    def test_split_node_splits_cp_metadata_by_owned_positions(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.page_size = 1

        root = TreeNode()
        root.key = RadixKey([])
        child = TreeNode()
        child.parent = root
        child.key = RadixKey(list(range(10)))
        child.value = None
        child.hash_value = None
        child.host_len = 10
        child.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=10,
            owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
            host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
            page_owners=torch.zeros(max(10, 0), dtype=torch.int8),
            page_size=1,
        )
        root.children[0] = child

        new_node = cache._split_node(child.key, child, 5)

        self.assertEqual(new_node.host_len, 5)
        self.assertEqual(child.host_len, 5)
        self.assertEqual(new_node.cp_hicache.owned_positions.tolist(), [0, 2])
        self.assertEqual(new_node.cp_hicache.host_indices.tolist(), [20, 21])
        self.assertEqual(child.cp_hicache.owned_positions.tolist(), [0, 4])
        self.assertEqual(child.cp_hicache.host_indices.tolist(), [22, 23])

    def test_cp_host_eviction_uses_physical_freed_slots_for_progress(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.root_node = TreeNode()
        cache.root_node.key = RadixKey([])
        cache.evictable_host_leaves = set()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.eviction_strategy = type(
            "Strategy", (), {"get_priority": lambda self, node: 0}
        )()
        cache._clear_pin = lambda node: None
        cache._record_remove_event = lambda node: None
        cache._update_host_leaf_status = lambda node: None
        cache._node_host_evict_indices = lambda node: torch.tensor(
            [99], dtype=torch.int64
        )
        freed = []
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "evict_host": lambda self, indices: freed.append(indices.clone())
                or len(indices)
            },
        )()
        node = TreeNode()
        node.parent = cache.root_node
        node.key = RadixKey([1, 2, 3, 4])
        node.value = None
        node.host_len = 4
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([1], dtype=torch.int64),
            host_indices=torch.tensor([70], dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        cache.root_node.children[1] = node
        cache.evictable_host_leaves.add(node)

        physical_freed = cache._evict_host_for_physical_slots(1)

        self.assertEqual(physical_freed, 1)
        self.assertEqual(freed[0].tolist(), [99])
        self.assertEqual(node.host_len, 0)
        self.assertIsNone(node.cp_hicache)

    def test_cp_host_eviction_unlinks_stale_leaf_to_free_parent_slots(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.root_node = TreeNode()
        cache.root_node.key = RadixKey([])
        cache.evictable_host_leaves = set()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.eviction_strategy = type(
            "Strategy", (), {"get_priority": lambda self, node: 0}
        )()
        cache._clear_pin = lambda node: None
        cache._record_remove_event = lambda node: None
        freed = []
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "evict_host": lambda self, indices: freed.append(indices.clone())
                or len(indices)
            },
        )()

        parent = TreeNode()
        parent.parent = cache.root_node
        parent.key = RadixKey([1])
        parent.value = None
        parent.host_len = 4
        parent.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([2], dtype=torch.int64),
            host_indices=torch.tensor([80], dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        cache.root_node.children[1] = parent

        stale_child = TreeNode()
        stale_child.parent = parent
        stale_child.key = RadixKey([2])
        stale_child.value = None
        stale_child.host_len = 0
        stale_child.cp_hicache = None
        parent.children[2] = stale_child
        cache.evictable_host_leaves.add(stale_child)

        physical_freed = cache._evict_host_for_physical_slots(1)

        self.assertEqual(physical_freed, 1)
        self.assertEqual(freed[0].tolist(), [80])
        self.assertNotIn(2, parent.children)
        self.assertEqual(parent.host_len, 0)
        self.assertIsNone(parent.cp_hicache)
        self.assertNotIn(1, cache.root_node.children)

    def test_cp_host_eviction_preserves_parent_with_sibling_after_stale_cleanup(
        self,
    ):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.root_node = TreeNode()
        cache.root_node.key = RadixKey([])
        cache.evictable_host_leaves = set()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.eviction_strategy = type(
            "Strategy", (), {"get_priority": lambda self, node: 0}
        )()
        cache._clear_pin = lambda node: None
        cache._record_remove_event = lambda node: None
        freed = []
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "evict_host": lambda self, indices: freed.append(indices.clone())
                or len(indices)
            },
        )()

        parent = TreeNode()
        parent.parent = cache.root_node
        parent.key = RadixKey([1])
        parent.value = None
        parent.host_len = 4
        parent.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([2], dtype=torch.int64),
            host_indices=torch.tensor([80], dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        cache.root_node.children[1] = parent

        stale_child = TreeNode()
        stale_child.parent = parent
        stale_child.key = RadixKey([2])
        stale_child.value = None
        stale_child.host_len = 0
        stale_child.cp_hicache = None
        parent.children[2] = stale_child

        sibling = TreeNode()
        sibling.parent = parent
        sibling.key = RadixKey([3])
        sibling.value = None
        sibling.host_len = 0
        sibling.cp_hicache = None
        parent.children[3] = sibling
        cache.evictable_host_leaves.add(stale_child)

        physical_freed = cache._evict_host_for_physical_slots(1)

        self.assertEqual(physical_freed, 0)
        self.assertEqual(freed, [])
        self.assertNotIn(2, parent.children)
        self.assertIn(1, cache.root_node.children)
        self.assertIs(cache.root_node.children[1], parent)
        self.assertIn(3, parent.children)
        self.assertIs(parent.children[3], sibling)
        self.assertEqual(parent.host_len, 4)
        self.assertIsNotNone(parent.cp_hicache)
        self.assertEqual(parent.cp_hicache.host_indices.tolist(), [80])

    def test_synchronized_cp_host_eviction_removes_zero_owned_logical_leaf(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.tp_world_size = 2
        cache.tp_group = object()
        cache._tp_group_rank = 0
        cache.root_node = TreeNode()
        cache.root_node.key = RadixKey([])
        cache.evictable_host_leaves = set()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.eviction_strategy = type(
            "Strategy", (), {"get_priority": lambda self, node: 0}
        )()
        cache._clear_pin = lambda node: None
        cache._record_remove_event = lambda node: None
        freed = []
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "evict_host": lambda self, indices: freed.append(indices.clone())
                or len(indices)
            },
        )()
        node = TreeNode()
        node.parent = cache.root_node
        node.key = RadixKey([1, 2, 3, 4])
        node.value = None
        node.host_len = 4
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.empty((0,), dtype=torch.int64),
            host_indices=torch.empty((0,), dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        cache.root_node.children[1] = node
        cache.evictable_host_leaves.add(node)

        all_done_states = iter([0, 1])

        def fake_all_reduce(done, op=None, group=None):
            done.fill_(next(all_done_states, 1))

        with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce):
            physical_freed = cache._evict_host_for_physical_slots(
                0, synchronize_across_ranks=True
            )

        self.assertEqual(physical_freed, 0)
        self.assertEqual(freed, [])
        self.assertNotIn(1, cache.root_node.children)
        self.assertEqual(node.host_len, 0)
        self.assertIsNone(node.cp_hicache)

    def test_cp_host_eviction_skips_all_reduce_without_tp_group(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.tp_world_size = 2
        cache.tp_group = None
        cache.root_node = TreeNode()
        cache.root_node.key = RadixKey([])
        cache.evictable_host_leaves = set()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.eviction_strategy = type(
            "Strategy", (), {"get_priority": lambda self, node: 0}
        )()
        cache._clear_pin = lambda node: None
        cache._record_remove_event = lambda node: None
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "evict_host": lambda self, indices: (_ for _ in ()).throw(
                    AssertionError("must not evict host slots")
                )
            },
        )()
        node = TreeNode()
        node.parent = cache.root_node
        node.key = RadixKey([1, 2, 3, 4])
        node.value = None
        node.host_len = 4
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.empty((0,), dtype=torch.int64),
            host_indices=torch.empty((0,), dtype=torch.int64),
            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
            page_size=1,
        )
        cache.root_node.children[1] = node
        cache.evictable_host_leaves.add(node)

        with patch(
            "torch.distributed.all_reduce",
            side_effect=AssertionError("must not all_reduce without tp_group"),
        ):
            physical_freed = cache._evict_host_for_physical_slots(
                0, synchronize_across_ranks=True
            )

        self.assertEqual(physical_freed, 0)
        self.assertIn(1, cache.root_node.children)
        self.assertEqual(node.host_len, 4)
        self.assertIsNotNone(node.cp_hicache)


class TestHiRadixCacheCPLoadBack(CustomTestCase):
    def test_cp_load_back_uses_host_len_not_host_value(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.root_node = TreeNode()
        cache.device = "cpu"
        cache.load_back_threshold = 1
        cache.evictable_size_ = 0
        cache.metrics_collector = None
        cache.ongoing_load_back = {}
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "load_cp": lambda self, nodes, node_id=-1: torch.arange(
                    32, 40, dtype=torch.int64
                )
            },
        )()
        cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
        cache.dec_lock_ref = lambda node: None
        cache.evict = lambda params: None
        parent = cache.root_node
        parent.key = RadixKey([])
        parent.value = torch.empty((0,), dtype=torch.int64)
        node = TreeNode()
        node.parent = parent
        node.key = RadixKey(list(range(8)))
        node.value = None
        node.host_value = None
        node.host_len = 8
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([0, 1], dtype=torch.int64),
            host_indices=torch.tensor([50, 51], dtype=torch.int64),
            page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
            page_size=1,
        )

        loaded = cache.load_back(node)

        self.assertEqual(loaded.tolist(), list(range(32, 40)))
        self.assertEqual(node.value.tolist(), list(range(32, 40)))

    def test_cp_load_back_threshold_uses_logical_length(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.root_node = TreeNode()
        cache.root_node.value = torch.empty((0,), dtype=torch.int64)
        cache.load_back_threshold = 5
        cache.evictable_size_ = 0
        cache.metrics_collector = None
        cache.ongoing_load_back = {}
        cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
        cache.dec_lock_ref = lambda node: None
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "load_cp": lambda self, nodes, node_id=-1: torch.arange(
                    10, 16, dtype=torch.int64
                )
            },
        )()
        node = TreeNode()
        node.parent = cache.root_node
        node.value = None
        node.host_len = 6
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=6,
            owned_positions=torch.empty((0,), dtype=torch.int64),
            host_indices=torch.empty((0,), dtype=torch.int64),
            page_owners=torch.zeros(max(6, 0), dtype=torch.int8),
            page_size=1,
        )

        loaded = cache.load_back(node)

        self.assertEqual(loaded.tolist(), [10, 11, 12, 13, 14, 15])

    def test_cp_match_prefix_counts_logical_host_hit_without_host_value(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.device = "cpu"
        cache.disable = False
        cache.page_size = 1
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.key_match_fn = lambda child_key, key: sum(
            1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
        )
        cache.maybe_bigram_convert = lambda key: (key, None)
        root = TreeNode()
        root.key = RadixKey([])
        root.value = torch.empty((0,), dtype=torch.int64)
        root.host_len = 0
        cache.root_node = root
        node = TreeNode()
        node.parent = root
        node.key = RadixKey(list(range(8)))
        node.value = None
        node.host_value = None
        node.host_len = 8
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([0, 1], dtype=torch.int64),
            host_indices=torch.tensor([50, 51], dtype=torch.int64),
            page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
            page_size=1,
        )
        root.children[0] = node

        result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8)))))

        self.assertEqual(result.host_hit_length, 8)
        self.assertIs(result.last_host_node, node)
        self.assertIs(result.last_device_node, root)

    def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.device = "cpu"
        cache.disable = False
        cache.page_size = 4
        cache.maybe_bigram_convert = lambda key: (key, None)
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.root_node = TreeNode()

        result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))

        self.assertEqual(result.device_indices.tolist(), [])
        self.assertIs(result.last_device_node, cache.root_node)
        self.assertIs(result.last_host_node, cache.root_node)
        self.assertEqual(result.host_hit_length, 0)

    def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = False
        cache.device = "cpu"
        cache.disable = False
        cache.page_size = 1
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache.key_match_fn = lambda child_key, key: sum(
            1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
        )
        cache.maybe_bigram_convert = lambda key: (key, None)
        root = TreeNode()
        root.key = RadixKey([])
        root.value = torch.empty((0,), dtype=torch.int64)
        root.host_value = None
        cache.root_node = root
        node = TreeNode()
        node.parent = root
        node.key = RadixKey(list(range(4)))
        node.value = torch.arange(4, dtype=torch.int64)
        node.host_value = None
        root.children[0] = node

        result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(4)))))

        self.assertEqual(result.host_hit_length, 0)
        self.assertIs(result.last_host_node, root)
        self.assertIs(result.last_device_node, node)

    def test_cp_load_back_non_evicted_node_returns_none_without_loading(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.load_back_threshold = 1
        cache.evictable_size_ = 0
        cache.metrics_collector = None
        cache.ongoing_load_back = {}
        cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
        cache.dec_lock_ref = lambda node: None
        cache.cache_controller = type(
            "Controller",
            (),
            {
                "load_cp": lambda self, nodes, node_id=-1: (_ for _ in ()).throw(
                    AssertionError("must not load without evicted nodes")
                )
            },
        )()
        node = TreeNode()
        node.value = torch.empty((0,), dtype=torch.int64)

        loaded = cache.load_back(node)

        self.assertIsNone(loaded)
        self.assertEqual(cache.ongoing_load_back, {})


if __name__ == "__main__":
    unittest.main()