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

import torch

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:
    kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio")
    for name in (
        "transfer_kv_all_layer",
        "transfer_kv_all_layer_direct_lf_pf",
        "transfer_kv_all_layer_lf_pf",
        "transfer_kv_all_layer_lf_ph",
        "transfer_kv_all_layer_mla",
        "transfer_kv_all_layer_mla_lf_pf",
        "transfer_kv_direct",
        "transfer_kv_per_layer",
        "transfer_kv_per_layer_direct_pf_lf",
        "transfer_kv_per_layer_mla",
        "transfer_kv_per_layer_mla_pf_lf",
        "transfer_kv_per_layer_pf_lf",
        "transfer_kv_per_layer_ph_lf",
    ):
        setattr(kvcacheio_stub, name, lambda *args, **kwargs: None)
    sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub

_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.managers.cache_controller import CacheOperation, HiCacheController
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
from sglang.srt.mem_cache.memory_pool_host import (
    HostKVCache,
    MHATokenToKVPoolHost,
    MLATokenToKVPoolHost,
    NSATokenToKVPoolHost,
    PreparedLoadDescriptor,
)
from sglang.srt.mem_cache.radix_cache import 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 FakeHostPool:
    def __init__(self, alloc_result):
        self.alloc_result = alloc_result
        self.alloc_calls = []
        self.backups = []
        self.layer_backups = []
        self.loads = []
        self.frees = []
        self.page_size = 4
        self.layout = "page_first_direct"

    def alloc(self, need_size):
        self.alloc_calls.append(need_size)
        if self.alloc_result is None:
            return None
        return self.alloc_result[:need_size].clone()

    def backup_from_device_all_layer(
        self, device_pool, host_indices, device_indices, io_backend
    ):
        self.backups.append((host_indices.clone(), device_indices.clone(), device_pool))

    def backup_from_device_per_layer(
        self, device_pool, host_indices, device_indices, layer_id, io_backend
    ):
        self.layer_backups.append(
            (host_indices.clone(), device_indices.clone(), layer_id, device_pool)
        )

    def load_to_device_per_layer(
        self, device_pool, host_indices, device_indices, layer_id, io_backend
    ):
        self.loads.append(
            (host_indices.clone(), device_indices.clone(), layer_id, device_pool)
        )

    def free(self, indices):
        self.frees.append(indices.clone())
        return len(indices)


class ContiguousPreferredHostPool(FakeHostPool):
    def __init__(self, alloc_result):
        super().__init__(alloc_result)
        self.contiguous_alloc_calls = []

    def alloc_contiguous_preferred(self, need_size):
        self.contiguous_alloc_calls.append(need_size)
        if self.alloc_result is None:
            return None
        return self.alloc_result[:need_size].clone()


class DummyHostKVCacheForAlloc(HostKVCache):
    def get_size_per_token(self):
        return 1

    def init_kv_buffer(self):
        return None

    def load_to_device_per_layer(
        self, device_pool, host_indices, device_indices, layer_id, io_backend
    ) -> None:
        pass

    def backup_from_device_per_layer(
        self, device_pool, host_indices, device_indices, layer_id, io_backend
    ) -> None:
        pass

    def backup_from_device_all_layer(
        self, device_pool, host_indices, device_indices, io_backend
    ) -> None:
        pass

    def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
        return torch.empty((0,), dtype=torch.uint8)

    def get_dummy_flat_data_page(self) -> torch.Tensor:
        return torch.empty((0,), dtype=torch.uint8)

    def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
        pass


class TestPreparedLoadDescriptor(CustomTestCase):
    def test_host_begin_load_builds_page_aligned_descriptor(self):
        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
        host_pool.page_size = 4
        host_pool.layout = "page_first_direct"

        host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
        device_indices = torch.tensor(
            [40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
        )

        host_pool.begin_load_to_device_op(
            host_indices, device_indices, io_backend="direct"
        )

        desc = host_pool._active_load_descriptor
        self.assertIsInstance(desc, PreparedLoadDescriptor)
        self.assertTrue(torch.equal(desc.host_indices, host_indices))
        self.assertTrue(torch.equal(desc.device_indices, device_indices))
        self.assertEqual(desc.num_tokens, 8)
        self.assertEqual(desc.num_pages, 2)
        self.assertEqual(desc.layout, "page_first_direct")
        self.assertEqual(desc.io_backend, "direct")
        self.assertEqual(desc.host_page_indices.tolist(), [2, 5])
        self.assertEqual(desc.device_page_indices.tolist(), [10, 16])

        host_pool.end_load_to_device_op()
        self.assertIsNone(host_pool._active_load_descriptor)

    def test_nsa_begin_load_attaches_indexer_pages_to_descriptor(self):
        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.page_size = 4
        host_pool.layout = "page_first_direct"
        host_pool.index_active_layer_ids = (0, 2)

        host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
        device_indices = torch.tensor(
            [40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
        )

        host_pool.begin_load_to_device_op(
            host_indices, device_indices, io_backend="direct"
        )

        desc = host_pool._active_load_descriptor
        self.assertEqual(desc.index_active_layer_ids, (0, 2))
        self.assertEqual(desc.index_host_page_indices.tolist(), [2, 5])
        self.assertEqual(desc.index_device_page_indices.tolist(), [10, 16])
        self.assertIs(host_pool._active_load_indexer_page_indices[0], desc.index_host_page_indices)
        self.assertIs(
            host_pool._active_load_indexer_page_indices[1],
            desc.index_device_page_indices,
        )

        host_pool.end_load_to_device_op()
        self.assertIsNone(host_pool._active_load_descriptor)
        self.assertIsNone(host_pool._active_load_indexer_page_indices)

    def test_missing_direct_load_descriptor_warns_once(self):
        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
        host_pool.page_size = 4
        host_pool.layout = "page_first_direct"

        with self.assertLogs(
            "sglang.srt.mem_cache.memory_pool_host", level="WARNING"
        ) as logs:
            self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
            self.assertIsNone(host_pool._get_active_load_descriptor("direct"))

        warnings = [
            line for line in logs.output if "missing_prepared_descriptor" in line
        ]
        self.assertEqual(len(warnings), 1)

    def test_mla_direct_load_uses_prepared_tai_descriptor_when_available(self):
        calls = []
        fake_desc = object()

        def fake_prepare(src_indices, dst_indices, **kwargs):
            calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
            return fake_desc

        def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
            calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))

        def fake_destroy(desc):
            calls.append(("destroy", desc))

        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with (
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
                return_value=fake_prepare,
            ),
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
                return_value=fake_submit,
            ),
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_destroy_h2d_page_descriptor",
                return_value=fake_destroy,
            ),
        ):
            host_pool.begin_load_to_device_op(
                host_indices, device_indices, io_backend="direct"
            )
            host_pool.load_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )
            host_pool.end_load_to_device_op()

        self.assertEqual(calls[0][0], "prepare")
        self.assertEqual(calls[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(calls[0][2].tolist(), [12, 13, 14, 15])
        self.assertEqual(calls[0][3], {"page_size": 4, "layout": "page_first_direct"})
        self.assertEqual(calls[1][0], "submit")
        self.assertIs(calls[1][1], fake_desc)
        self.assertEqual(calls[1][2][0].data_ptr(), host_pool.kv_buffer.data_ptr())
        self.assertEqual(
            calls[1][3][0].data_ptr(), device_pool.kv_buffer[2].data_ptr()
        )
        self.assertEqual(calls[1][4], {"layer_id": 2})
        self.assertEqual(calls[2], ("destroy", fake_desc))

    def test_nsa_index_direct_load_uses_prepared_tai_index_descriptor(self):
        calls = []
        fake_base_desc = object()
        fake_index_desc = object()

        def fake_prepare(src_indices, dst_indices, **kwargs):
            calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
            if kwargs["page_size"] == 1:
                return fake_index_desc
            return fake_base_desc

        def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
            calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.index_active_layer_ids = (0, 1, 2)
        host_pool.index_k_with_scale_buffer = torch.empty(
            (8, 3, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with (
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
                return_value=fake_prepare,
            ),
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
                return_value=fake_submit,
            ),
        ):
            host_pool.begin_load_to_device_op(
                host_indices, device_indices, io_backend="direct"
            )
            host_pool._load_indexer_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=1,
                io_backend="direct",
            )

        self.assertEqual([c[0] for c in calls[:2]], ["prepare", "prepare"])
        self.assertEqual(calls[1][1].tolist(), [1])
        self.assertEqual(calls[1][2].tolist(), [3])
        self.assertEqual(calls[1][3], {"page_size": 1, "layout": "page_first_direct"})
        self.assertEqual(calls[2][0], "submit")
        self.assertIs(calls[2][1], fake_index_desc)
        self.assertEqual(
            calls[2][2][0].data_ptr(), host_pool.index_k_with_scale_buffer.data_ptr()
        )
        self.assertEqual(
            calls[2][3][0].data_ptr(),
            device_pool.index_k_with_scale_buffer[1].data_ptr(),
        )
        self.assertEqual(calls[2][4], {"layer_id": 1})


class FakeDevicePool:
    device = "cpu"
    layer_num = 1

    def __init__(self, name="target", layer_num=1):
        self.name = name
        self.layer_num = layer_num
        self.layer_backup_notifiers = []

    def register_layer_transfer_counter(self, counter):
        self.counter = counter

    def register_layer_backup_notifier(self, notifier):
        self.layer_backup_notifiers.append(notifier)

    def notify_layer_end_for_backup(self, layer_id):
        for notifier in self.layer_backup_notifiers:
            notifier(layer_id)


class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
    def test_mla_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self):
        calls = []

        def fake_kernel(src, dst, src_indices, dst_indices, **kwargs):
            calls.append((src, dst, src_indices.clone(), dst_indices.clone(), kwargs))

        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
        host_pool.layout = "page_first"
        host_pool.token_stride_size = 16
        host_pool.layout_dim = 64
        host_pool.kv_buffer = torch.empty((32, 4, 1, 16), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.kv_buffer = torch.empty((4, 32, 1, 16), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_mla_lf_pf",
            return_value=fake_kernel,
        ):
            host_pool.backup_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="kernel",
            )

        self.assertEqual(len(calls), 1)
        src, dst, src_indices, dst_indices, kwargs = calls[0]
        expected_src = device_pool.kv_buffer[2]
        self.assertEqual(src.data_ptr(), expected_src.data_ptr())
        self.assertEqual(src.shape, expected_src.shape)
        self.assertEqual(src.stride(), expected_src.stride())
        self.assertIs(dst, host_pool.kv_buffer)
        self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["item_size"], 16)
        self.assertEqual(kwargs["dst_layout_dim"], 64)

    def test_nsa_indexer_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self):
        calls = []

        def fake_kernel(src, dst, src_indices, dst_indices, **kwargs):
            calls.append((src, dst, src_indices.clone(), dst_indices.clone(), kwargs))

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "page_first"
        host_pool.page_size = 4
        host_pool.indexer_page_stride_size = 32
        host_pool.indexer_layout_dim = 96
        host_pool.index_k_with_scale_buffer = torch.empty(
            (16, 3, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 16, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
        device_indices = torch.tensor(
            [12, 13, 14, 15, 28, 29, 30, 31], dtype=torch.int64
        )

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_mla_lf_pf",
            return_value=fake_kernel,
        ):
            host_pool._backup_indexer_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=1,
                io_backend="kernel",
            )

        self.assertEqual(len(calls), 1)
        src, dst, src_indices, dst_indices, kwargs = calls[0]
        expected_src = device_pool.index_k_with_scale_buffer[1]
        self.assertEqual(src.data_ptr(), expected_src.data_ptr())
        self.assertEqual(src.shape, expected_src.shape)
        self.assertEqual(src.stride(), expected_src.stride())
        self.assertIs(dst, host_pool.index_k_with_scale_buffer)
        self.assertEqual(src_indices.tolist(), [3, 7])
        self.assertEqual(dst_indices.tolist(), [2, 5])
        self.assertEqual(kwargs["layer_id"], 1)
        self.assertEqual(kwargs["item_size"], 32)
        self.assertEqual(kwargs["dst_layout_dim"], 96)

    def test_mla_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
            return_value=fake_direct,
        ):
            host_pool.backup_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(src_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
        self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_mha_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((2, 8, 3, 4, 2, 8), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        device_pool.v_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
            return_value=fake_direct,
        ):
            host_pool.backup_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 2)
        self.assertEqual(src_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
        self.assertEqual(src_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
        self.assertEqual(len(dst_ptrs), 2)
        self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
        self.assertEqual(dst_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
        self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_nsa_indexer_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.index_k_with_scale_buffer = torch.empty(
            (8, 3, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
            return_value=fake_direct,
        ):
            host_pool._backup_indexer_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=1,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(
            src_ptrs[0].data_ptr(),
            device_pool.index_k_with_scale_buffer[1].data_ptr(),
        )
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(
            dst_ptrs[0].data_ptr(),
            host_pool.index_k_with_scale_buffer.data_ptr(),
        )
        self.assertEqual(src_indices.tolist(), [3])
        self.assertEqual(dst_indices.tolist(), [1])
        self.assertEqual(kwargs["layer_id"], 1)
        self.assertEqual(kwargs["page_size"], 1)

    def test_mla_layer_page_first_per_layer_backup_uses_direct_lf_lpf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
        host_pool.layout = "layer_page_first"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((3, 8, 4, 1, 16), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
            return_value=fake_direct,
        ):
            host_pool.backup_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(src_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
        self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_mha_layer_page_first_per_layer_backup_uses_direct_lf_lpf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
        host_pool.layout = "layer_page_first"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((2, 3, 8, 4, 2, 8), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        device_pool.v_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
            return_value=fake_direct,
        ):
            host_pool.backup_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 2)
        self.assertEqual(src_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
        self.assertEqual(src_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
        self.assertEqual(len(dst_ptrs), 2)
        self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
        self.assertEqual(dst_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
        self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_nsa_indexer_layer_page_first_per_layer_backup_uses_direct_lf_lpf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "layer_page_first"
        host_pool.page_size = 4
        host_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
            return_value=fake_direct,
        ):
            host_pool._backup_indexer_from_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=1,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(
            src_ptrs[0].data_ptr(),
            device_pool.index_k_with_scale_buffer[1].data_ptr(),
        )
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(
            dst_ptrs[0].data_ptr(),
            host_pool.index_k_with_scale_buffer.data_ptr(),
        )
        self.assertEqual(src_indices.tolist(), [3])
        self.assertEqual(dst_indices.tolist(), [1])
        self.assertEqual(kwargs["layer_id"], 1)
        self.assertEqual(kwargs["page_size"], 1)

    def test_mla_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
            return_value=fake_direct,
        ):
            host_pool.load_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(src_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
        self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_mha_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((2, 8, 3, 4, 2, 8), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        device_pool.v_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
            return_value=fake_direct,
        ):
            host_pool.load_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 2)
        self.assertEqual(src_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
        self.assertEqual(src_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
        self.assertEqual(len(dst_ptrs), 2)
        self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
        self.assertEqual(dst_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
        self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_nsa_indexer_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.index_k_with_scale_buffer = torch.empty(
            (8, 3, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
            return_value=fake_direct,
        ):
            host_pool._load_indexer_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=1,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(
            src_ptrs[0].data_ptr(),
            host_pool.index_k_with_scale_buffer.data_ptr(),
        )
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(
            dst_ptrs[0].data_ptr(),
            device_pool.index_k_with_scale_buffer[1].data_ptr(),
        )
        self.assertEqual(src_indices.tolist(), [1])
        self.assertEqual(dst_indices.tolist(), [3])
        self.assertEqual(kwargs["layer_id"], 1)
        self.assertEqual(kwargs["page_size"], 1)

    def test_mla_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
        host_pool.layout = "layer_page_first"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((3, 8, 4, 1, 16), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lpf_lf",
            return_value=fake_direct,
        ):
            host_pool.load_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(src_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
        self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_mha_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
        host_pool.layout = "layer_page_first"
        host_pool.page_size = 4
        host_pool.kv_buffer = torch.empty((2, 3, 8, 4, 2, 8), dtype=torch.uint8)
        device_pool = type("DevicePool", (), {})()
        device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        device_pool.v_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lpf_lf",
            return_value=fake_direct,
        ):
            host_pool.load_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=2,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 2)
        self.assertEqual(src_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
        self.assertEqual(src_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
        self.assertEqual(len(dst_ptrs), 2)
        self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
        self.assertEqual(dst_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
        self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
        self.assertEqual(kwargs["layer_id"], 2)
        self.assertEqual(kwargs["page_size"], 4)

    def test_nsa_indexer_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "layer_page_first"
        host_pool.page_size = 4
        host_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        with patch(
            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lpf_lf",
            return_value=fake_direct,
        ):
            host_pool._load_indexer_to_device_per_layer(
                device_pool,
                host_indices,
                device_indices,
                layer_id=1,
                io_backend="direct",
            )

        self.assertEqual(len(calls), 1)
        src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
        self.assertEqual(len(src_ptrs), 1)
        self.assertEqual(
            src_ptrs[0].data_ptr(),
            host_pool.index_k_with_scale_buffer.data_ptr(),
        )
        self.assertEqual(len(dst_ptrs), 1)
        self.assertEqual(
            dst_ptrs[0].data_ptr(),
            device_pool.index_k_with_scale_buffer[1].data_ptr(),
        )
        self.assertEqual(src_indices.tolist(), [1])
        self.assertEqual(dst_indices.tolist(), [3])
        self.assertEqual(kwargs["layer_id"], 1)
        self.assertEqual(kwargs["page_size"], 1)

    def test_nsa_indexer_load_reuses_precomputed_page_indices_across_layers(self):
        calls = []

        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
            calls.append(
                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
            )

        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
        host_pool.layout = "page_first_direct"
        host_pool.page_size = 4
        host_pool.index_k_with_scale_buffer = torch.empty(
            (8, 3, 1, 32), dtype=torch.uint8
        )
        device_pool = type("DevicePool", (), {})()
        device_pool.index_k_with_scale_buffer = torch.empty(
            (3, 8, 32), dtype=torch.uint8
        )
        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)

        original_getter = host_pool._get_indexer_page_indices
        getter_calls = []

        def counting_getter(h, d):
            getter_calls.append((h.clone(), d.clone()))
            return original_getter(h, d)

        host_pool._get_indexer_page_indices = counting_getter

        with (
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
                side_effect=RuntimeError("missing prepared descriptor api"),
            ),
            patch(
                "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
                return_value=fake_direct,
            ),
        ):
            host_pool.begin_load_to_device_op(
                host_indices, device_indices, io_backend="direct"
            )
            try:
                for layer_id in range(3):
                    host_pool._load_indexer_to_device_per_layer(
                        device_pool,
                        host_indices,
                        device_indices,
                        layer_id=layer_id,
                        io_backend="direct",
                    )
            finally:
                host_pool.end_load_to_device_op()

        self.assertEqual(len(getter_calls), 1)
        self.assertEqual(len(calls), 3)
        self.assertEqual([call[4]["layer_id"] for call in calls], [0, 1, 2])
        self.assertEqual([call[2].tolist() for call in calls], [[1], [1], [1]])
        self.assertEqual([call[3].tolist() for call in calls], [[3], [3], [3]])


class FakeAllocator:
    def __init__(self, alloc_result=None):
        self.alloc_result = alloc_result
        self.alloc_calls = []
        self.owner_alloc_calls = []
        self.frees = []
        self.cp_size = 4
        self.cp_rank = 1
        self.page_size = 4
        self.device_pool = FakeDevicePool()

    def get_kvcache(self):
        return self.device_pool

    def alloc(self, need_size):
        self.alloc_calls.append(need_size)
        if self.alloc_result is None:
            return None
        return self.alloc_result[:need_size].clone()

    def alloc_pages_with_owners(self, page_owners):
        owners = list(page_owners)
        self.owner_alloc_calls.append(owners)
        if self.alloc_result is None:
            return None
        need_size = len(owners) * self.page_size
        return self.alloc_result[:need_size].clone()

    def free(self, indices):
        self.frees.append(indices.clone())
        return len(indices)


class HostIndicesTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data):
        return torch.Tensor._make_subclass(cls, data, require_grad=False)

    def to(self, *args, **kwargs):
        raise AssertionError("load_cp should not move host indices before queuing")


class DummyEvent:
    def record(self):
        pass

    def wait(self, stream):
        pass

    def query(self):
        return True

    def synchronize(self):
        pass


class DummyStream:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class DummyDeviceModule:
    Event = DummyEvent
    Stream = DummyStream

    @staticmethod
    def stream(stream):
        return stream


class DummyLayerDoneCounter:
    def __init__(self):
        self.events = [
            type(
                "ProducerEvent",
                (),
                {
                    "start_event": DummyEvent(),
                    "finish_event": DummyEvent(),
                    "complete": lambda self, layer_id: None,
                },
            )()
        ]

    def update_producer(self):
        return 0


class RecordingProducerEvent:
    def __init__(self, order):
        self.start_event = DummyEvent()
        self.finish_event = DummyEvent()
        self.order = order

    def complete(self, layer_id):
        self.order.append(("complete", layer_id))


class RecordingLayerDoneCounter:
    def __init__(self, order):
        self.events = [RecordingProducerEvent(order)]

    def update_producer(self):
        return 0


class TestHiCacheControllerCPWrite(CustomTestCase):
    def setUp(self):
        self.device_module_patcher = patch(
            "sglang.srt.managers.cache_controller.device_module",
            DummyDeviceModule,
        )
        self.nsa_pool_patcher = patch(
            "sglang.srt.managers.cache_controller.NSATokenToKVPool",
            FakeDevicePool,
        )
        self.device_module_patcher.start()
        self.nsa_pool_patcher.start()
        self.addCleanup(self.device_module_patcher.stop)
        self.addCleanup(self.nsa_pool_patcher.stop)

    def make_controller(
        self,
        host_pool,
        allocator=None,
        cp_rank=1,
        draft_host_pool=None,
        draft_mem_pool_device=None,
    ):
        allocator = allocator or FakeAllocator()
        controller = HiCacheController(
            token_to_kv_pool_allocator=allocator,
            mem_pool_host=host_pool,
            page_size=4,
            tp_group=None,
            load_cache_event=__import__("threading").Event(),
            io_backend="direct",
            cp_shared_kv_layout=CpSharedKVLayout(
                page_size=4, cp_size=4, cp_rank=cp_rank
            ),
            draft_mem_pool_host=draft_host_pool,
            draft_mem_pool_device=draft_mem_pool_device,
        )
        controller.layer_done_counter = DummyLayerDoneCounter()
        return controller

    def test_cp_write_filters_to_owned_physical_locs(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=7)

        self.assertEqual(result.metadata.logical_len, 16)
        self.assertEqual(result.metadata.owned_positions.tolist(), [4, 5, 6, 7])
        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(host_pool.alloc_calls, [4])
        self.assertEqual(host_pool.backups, [])
        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])

    def test_cp_write_accepts_valid_tail_and_pads_owned_physical_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)

        result = controller.write(logical_locs, node_id=21)

        self.assertEqual(result.metadata.logical_len, 3)
        self.assertEqual(result.metadata.valid_len, 3)
        self.assertEqual(result.metadata.padded_len, 4)
        self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3])
        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(host_pool.alloc_calls, [4])
        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])

    def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.tensor([8, 9, 11, 10], dtype=torch.int64)

        with self.assertRaisesRegex(
            ValueError, "physical_device_indices.*contiguous page spans"
        ):
            controller.write(logical_locs, node_id=22)

    def test_cp_write_rejects_non_contiguous_host_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 103, 102], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(8, 12, dtype=torch.int64)

        with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
            controller.write(logical_locs, node_id=23)

    def test_cp_write_zero_owned_returns_metadata_and_noop_ack(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=3)
        logical_locs = torch.arange(4, 8, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=8)

        self.assertEqual(result.metadata.logical_len, 4)
        self.assertEqual(result.metadata.host_indices.tolist(), [])
        self.assertEqual(host_pool.alloc_calls, [])
        self.assertEqual(len(controller.ack_write_queue), 1)

    def test_cp_write_zero_owned_with_draft_returns_empty_draft_metadata(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            cp_rank=3,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        logical_locs = torch.arange(4, 8, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=18)

        self.assertEqual(result.metadata.logical_len, 4)
        self.assertEqual(result.metadata.host_indices.tolist(), [])
        self.assertEqual(result.metadata.draft_host_indices.tolist(), [])
        self.assertEqual(host_pool.alloc_calls, [])
        self.assertEqual(draft_host_pool.alloc_calls, [])
        self.assertEqual(len(controller.ack_write_queue), 1)

    def test_cp_write_allocation_failure_reports_required_host_slots(self):
        host_pool = FakeHostPool(None)
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=9)

        self.assertEqual(result.required_host_slots, 4)

    def test_cp_write_with_draft_pool_backs_target_and_draft_locs(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        draft_device_pool = FakeDevicePool("draft")
        controller = self.make_controller(
            host_pool,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=77)

        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(
            result.metadata.draft_host_indices.tolist(), [200, 201, 202, 203]
        )
        self.assertEqual(host_pool.alloc_calls, [4])
        self.assertEqual(draft_host_pool.alloc_calls, [4])
        self.assertEqual(host_pool.backups, [])
        self.assertEqual(draft_host_pool.backups, [])
        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertIs(draft_host_pool.layer_backups[0][3], draft_device_pool)
        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [77])

    def test_cp_write_valid_tail_with_draft_mirrors_target_padded_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        draft_device_pool = FakeDevicePool("draft")
        controller = self.make_controller(
            host_pool,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)

        result = controller.write(logical_locs, node_id=177)

        self.assertEqual(result.metadata.logical_len, 3)
        self.assertEqual(result.metadata.padded_len, 4)
        self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3])
        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(
            result.metadata.draft_host_indices.tolist(), [200, 201, 202, 203]
        )
        self.assertEqual(host_pool.alloc_calls, [4])
        self.assertEqual(draft_host_pool.alloc_calls, [4])
        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertIs(draft_host_pool.layer_backups[0][3], draft_device_pool)

    def test_cp_write_draft_allocation_failure_rolls_back_target_host(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(None)
        controller = self.make_controller(
            host_pool,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=78)

        self.assertIsNone(result.metadata)
        self.assertEqual(result.required_host_slots, 4)
        self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
        self.assertEqual(host_pool.backups, [])
        self.assertEqual(draft_host_pool.backups, [])

    def test_cp_reserve_write_queues_no_transfer_until_submit(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        reservation = controller.reserve_write_cp(logical_locs, node_id=79)

        self.assertEqual(reservation.metadata.logical_len, 16)
        self.assertEqual(host_pool.alloc_calls, [4])
        self.assertEqual(host_pool.backups, [])
        self.assertEqual(controller.write_queue, [])
        self.assertEqual(controller.ack_write_queue, [])

        with self.assertLogs(
            "sglang.srt.managers.cache_controller", level="WARNING"
        ) as logs:
            controller.submit_write_cp_all_layer(reservation)

        self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [79])
        self.assertIn("all-layer backup fallback", "\n".join(logs.output))

    def test_cp_reserve_write_uses_contiguous_preferred_host_alloc(self):
        host_pool = ContiguousPreferredHostPool(
            torch.tensor([100, 101, 102, 103], dtype=torch.int64)
        )
        draft_host_pool = ContiguousPreferredHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        controller = self.make_controller(
            host_pool,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        reservation = controller.reserve_write_cp(logical_locs, node_id=179)

        self.assertEqual(reservation.metadata.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(
            reservation.metadata.draft_host_indices.tolist(), [200, 201, 202, 203]
        )
        self.assertEqual(host_pool.contiguous_alloc_calls, [4])
        self.assertEqual(draft_host_pool.contiguous_alloc_calls, [4])
        self.assertEqual(host_pool.alloc_calls, [])
        self.assertEqual(draft_host_pool.alloc_calls, [])

    def test_host_alloc_contiguous_preferred_skips_fragmented_fifo_prefix(self):
        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
        host_pool.page_size = 4
        host_pool.lock = __import__("threading").RLock()
        host_pool.free_slots = torch.tensor(
            [100, 101, 102, 103, 8, 9, 10, 11, 12, 13, 14, 15],
            dtype=torch.int64,
        )

        selected = host_pool.alloc_contiguous_preferred(8)

        self.assertEqual(selected.tolist(), [8, 9, 10, 11, 12, 13, 14, 15])
        self.assertEqual(host_pool.free_slots.tolist(), [100, 101, 102, 103])

    def test_host_alloc_contiguous_preferred_uses_lazy_extent_index(self):
        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
        host_pool.page_size = 4
        host_pool.lock = __import__("threading").RLock()
        pages = [50, 51, 52, 53, 100, 7, 8]
        host_pool.free_slots = torch.tensor(
            [page * 4 + offset for page in pages for offset in range(4)],
            dtype=torch.int64,
        )

        selected = host_pool.alloc_contiguous_preferred(16)

        self.assertEqual(
            selected.tolist(),
            [page * 4 + offset for page in [50, 51, 52, 53] for offset in range(4)],
        )
        self.assertEqual(host_pool.available_size(), 12)
        self.assertTrue(host_pool._free_slots_dirty)
        self.assertEqual(
            host_pool.free_slots.tolist(),
            [page * 4 + offset for page in [7, 8, 100] for offset in range(4)],
        )

    def test_cp_reserve_zero_owned_queues_no_ack_until_submit(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=3)
        logical_locs = torch.arange(4, 8, dtype=torch.int64)

        reservation = controller.reserve_write_cp(logical_locs, node_id=80)

        self.assertEqual(reservation.metadata.host_indices.tolist(), [])
        self.assertEqual(host_pool.alloc_calls, [])
        self.assertEqual(controller.ack_write_queue, [])

        with self.assertLogs(
            "sglang.srt.managers.cache_controller", level="WARNING"
        ) as logs:
            controller.submit_write_cp_all_layer(reservation)

        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [80])
        self.assertEqual(host_pool.backups, [])
        self.assertIn("all-layer backup fallback", "\n".join(logs.output))

    def test_cp_reserve_draft_allocation_failure_rolls_back_without_transfer(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(None)
        controller = self.make_controller(
            host_pool,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.reserve_write_cp(logical_locs, node_id=81)

        self.assertIsNone(result.metadata)
        self.assertEqual(result.required_host_slots, 4)
        self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
        self.assertEqual(host_pool.backups, [])
        self.assertEqual(draft_host_pool.backups, [])
        self.assertEqual(controller.write_queue, [])
        self.assertEqual(controller.draft_write_queue, [])

    def test_cp_submit_write_cp_layer_pairs_target_and_draft_with_single_final_ack(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=2)
        draft_device_pool = FakeDevicePool("draft", layer_num=2)
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        logical_locs = torch.arange(4, 20, dtype=torch.int64)
        reservation = controller.reserve_write_cp(logical_locs, node_id=82)

        controller.submit_write_cp_layer(reservation, 0)

        self.assertEqual(controller.ack_write_queue, [])
        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(host_pool.layer_backups[0][2], 0)
        self.assertEqual(draft_host_pool.layer_backups[0][2], 0)

        controller.submit_write_cp_layer(reservation, 1)

        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [82])
        self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1])
        self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0, 1])

    def test_cp_submit_write_cp_layer_zero_owned_final_ack_once(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=2)
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3)
        logical_locs = torch.arange(4, 8, dtype=torch.int64)
        reservation = controller.reserve_write_cp(logical_locs, node_id=83)

        controller.submit_write_cp_layer(reservation, 0)
        self.assertEqual(controller.ack_write_queue, [])

        controller.submit_write_cp_layer(reservation, 1)

        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [83])
        self.assertEqual(host_pool.layer_backups, [])

    def test_cp_layer_hook_submits_registered_write_without_all_layer_backup(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=2)
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        logical_locs = torch.arange(4, 20, dtype=torch.int64)
        reservation = controller.reserve_write_cp(logical_locs, node_id=84)

        controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False)

        self.assertEqual(host_pool.backups, [])
        self.assertEqual(host_pool.layer_backups, [])
        self.assertEqual(controller.ack_write_queue, [])

        allocator.device_pool.notify_layer_end_for_backup(0)

        self.assertEqual(host_pool.layer_backups[0][2], 0)
        self.assertEqual(controller.ack_write_queue, [])

        allocator.device_pool.notify_layer_end_for_backup(1)

        self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1])
        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [84])
        self.assertEqual(host_pool.backups, [])

    def test_cp_layer_hook_groups_target_backups_across_pending_reservations(self):
        class SequentialHostPool(FakeHostPool):
            def __init__(self, alloc_results):
                super().__init__(torch.empty((0,), dtype=torch.int64))
                self.alloc_results = [result.clone() for result in alloc_results]

            def alloc(self, need_size):
                self.alloc_calls.append(need_size)
                if not self.alloc_results:
                    return None
                return self.alloc_results.pop(0).clone()

        host_pool = SequentialHostPool(
            [
                torch.tensor([100, 101, 102, 103], dtype=torch.int64),
                torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            ]
        )
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=2)
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        reservation_a = controller.reserve_write_cp(
            torch.arange(4, 20, dtype=torch.int64), node_id=501
        )
        reservation_b = controller.reserve_write_cp(
            torch.arange(20, 36, dtype=torch.int64), node_id=502
        )
        controller.submit_write_cp_per_layer(reservation_a, catch_up_all_layers=False)
        controller.submit_write_cp_per_layer(reservation_b, catch_up_all_layers=False)

        allocator.device_pool.notify_layer_end_for_backup(0)

        self.assertEqual(len(host_pool.layer_backups), 1)
        host_indices, device_indices, layer_id, device_pool = host_pool.layer_backups[0]
        self.assertEqual(
            host_indices.tolist(), [100, 101, 102, 103, 200, 201, 202, 203]
        )
        self.assertEqual(device_indices.tolist(), [4, 5, 6, 7, 8, 9, 10, 11])
        self.assertEqual(layer_id, 0)
        self.assertIs(device_pool, allocator.device_pool)
        self.assertEqual(controller.ack_write_queue, [])

        allocator.device_pool.notify_layer_end_for_backup(1)

        self.assertEqual(len(host_pool.layer_backups), 2)
        self.assertEqual([backup[2] for backup in host_pool.layer_backups], [0, 1])
        self.assertEqual(len(controller.ack_write_queue), 2)
        self.assertEqual(
            [ack.node_ids for ack in controller.ack_write_queue], [[501], [502]]
        )

    def test_cp_layer_hook_groups_target_and_draft_backups_by_source(self):
        class SequentialHostPool(FakeHostPool):
            def __init__(self, alloc_results):
                super().__init__(torch.empty((0,), dtype=torch.int64))
                self.alloc_results = [result.clone() for result in alloc_results]

            def alloc(self, need_size):
                self.alloc_calls.append(need_size)
                if not self.alloc_results:
                    return None
                return self.alloc_results.pop(0).clone()

        host_pool = SequentialHostPool(
            [
                torch.tensor([100, 101, 102, 103], dtype=torch.int64),
                torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            ]
        )
        draft_host_pool = SequentialHostPool(
            [
                torch.tensor([300, 301, 302, 303], dtype=torch.int64),
                torch.tensor([400, 401, 402, 403], dtype=torch.int64),
            ]
        )
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=2)
        draft_device_pool = FakeDevicePool("draft", layer_num=2)
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        reservation_a = controller.reserve_write_cp(
            torch.arange(4, 20, dtype=torch.int64), node_id=601
        )
        reservation_b = controller.reserve_write_cp(
            torch.arange(20, 36, dtype=torch.int64), node_id=602
        )
        controller.submit_write_cp_per_layer(reservation_a, catch_up_all_layers=False)
        controller.submit_write_cp_per_layer(reservation_b, catch_up_all_layers=False)

        allocator.device_pool.notify_layer_end_for_backup(0)

        self.assertEqual(len(host_pool.layer_backups), 1)
        self.assertEqual(len(draft_host_pool.layer_backups), 0)
        self.assertEqual(
            host_pool.layer_backups[0][0].tolist(),
            [100, 101, 102, 103, 200, 201, 202, 203],
        )
        self.assertEqual(
            host_pool.layer_backups[0][1].tolist(),
            [4, 5, 6, 7, 8, 9, 10, 11],
        )

        draft_device_pool.notify_layer_end_for_backup(0)

        self.assertEqual(len(draft_host_pool.layer_backups), 1)
        self.assertEqual(
            draft_host_pool.layer_backups[0][0].tolist(),
            [300, 301, 302, 303, 400, 401, 402, 403],
        )
        self.assertEqual(
            draft_host_pool.layer_backups[0][1].tolist(),
            [4, 5, 6, 7, 8, 9, 10, 11],
        )
        self.assertEqual(controller.ack_write_queue, [])

        allocator.device_pool.notify_layer_end_for_backup(1)
        self.assertEqual(len(host_pool.layer_backups), 2)
        self.assertEqual(controller.ack_write_queue, [])

        draft_device_pool.notify_layer_end_for_backup(1)

        self.assertEqual(len(draft_host_pool.layer_backups), 2)
        self.assertEqual(len(controller.ack_write_queue), 2)
        self.assertEqual(
            [ack.node_ids for ack in controller.ack_write_queue], [[601], [602]]
        )

    def test_cp_layer_hook_keeps_zero_owned_ack_in_grouped_backup(self):
        class SequentialHostPool(FakeHostPool):
            def __init__(self, alloc_results):
                super().__init__(torch.empty((0,), dtype=torch.int64))
                self.alloc_results = [result.clone() for result in alloc_results]

            def alloc(self, need_size):
                self.alloc_calls.append(need_size)
                if not self.alloc_results:
                    return None
                return self.alloc_results.pop(0).clone()

        host_pool = SequentialHostPool(
            [torch.tensor([100, 101, 102, 103], dtype=torch.int64)]
        )
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=2)
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        zero_owned = controller.reserve_write_cp(
            torch.arange(4, 8, dtype=torch.int64), node_id=701
        )
        owned = controller.reserve_write_cp(
            torch.arange(8, 24, dtype=torch.int64), node_id=702
        )
        controller.submit_write_cp_per_layer(zero_owned, catch_up_all_layers=False)
        controller.submit_write_cp_per_layer(owned, catch_up_all_layers=False)

        allocator.device_pool.notify_layer_end_for_backup(0)
        allocator.device_pool.notify_layer_end_for_backup(1)

        self.assertEqual(len(host_pool.layer_backups), 2)
        for host_indices, device_indices, _, _ in host_pool.layer_backups:
            self.assertEqual(host_indices.tolist(), [100, 101, 102, 103])
            self.assertEqual(device_indices.tolist(), [4, 5, 6, 7])
        self.assertEqual(len(controller.ack_write_queue), 2)
        self.assertEqual(
            [ack.node_ids for ack in controller.ack_write_queue], [[701], [702]]
        )

    def test_cp_layer_hook_waits_for_draft_source_before_final_ack(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        allocator = FakeAllocator()
        allocator.device_pool = FakeDevicePool("target", layer_num=1)
        draft_device_pool = FakeDevicePool("draft", layer_num=1)
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        logical_locs = torch.arange(4, 20, dtype=torch.int64)
        reservation = controller.reserve_write_cp(logical_locs, node_id=85)

        controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False)
        allocator.device_pool.notify_layer_end_for_backup(0)

        self.assertEqual([x[2] for x in host_pool.layer_backups], [0])
        self.assertEqual(draft_host_pool.layer_backups, [])
        self.assertEqual(controller.ack_write_queue, [])

        draft_device_pool.notify_layer_end_for_backup(0)

        self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0])
        self.assertEqual(len(controller.ack_write_queue), 1)
        self.assertEqual(controller.ack_write_queue[0].node_ids, [85])

    def test_generate_storage_config_constructs_config_at_runtime(self):
        controller = HiCacheController.__new__(HiCacheController)
        controller.mem_pool_device = FakeDevicePool()
        controller.mem_pool_host = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller.pp_rank = 1
        controller.pp_size = 2
        controller.enable_storage_metrics = True

        with patch(
            "sglang.srt.managers.cache_controller.is_dp_attention_enabled",
            return_value=False,
        ), patch(
            "sglang.srt.managers.cache_controller.get_tensor_model_parallel_rank",
            return_value=3,
        ), patch(
            "sglang.srt.managers.cache_controller.get_tensor_model_parallel_world_size",
            return_value=4,
        ):
            config = controller._generate_storage_config(
                model_name="test-model",
                storage_backend_extra_config={"tp_lcm_size": 8},
            )

        self.assertEqual(config.tp_rank, 3)
        self.assertEqual(config.tp_size, 4)
        self.assertEqual(config.pp_rank, 1)
        self.assertEqual(config.pp_size, 2)
        self.assertEqual(config.model_name, "test-model")
        self.assertEqual(config.tp_lcm_size, 8)

    def test_attach_storage_backend_rejects_cp_hicache(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(host_pool)

        with self.assertRaisesRegex(RuntimeError, "CP shared KV.*storage backend"):
            controller.attach_storage_backend("mooncake")


class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
    def test_cp_start_loading_batches_multiple_load_cp_requests_with_draft(self):
        class SequentialOwnerAllocator(FakeAllocator):
            def __init__(self, alloc_results):
                super().__init__()
                self.alloc_results = [result.clone() for result in alloc_results]
                self.device_pool = FakeDevicePool("target", layer_num=2)

            def alloc_pages_with_owners(self, page_owners):
                owners = list(page_owners)
                self.owner_alloc_calls.append(owners)
                if not self.alloc_results:
                    return None
                return self.alloc_results.pop(0).clone()

        host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
        draft_host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
        allocator = SequentialOwnerAllocator(
            [
                torch.arange(64, 80, dtype=torch.int64),
                torch.arange(80, 96, dtype=torch.int64),
            ]
        )
        draft_device_pool = FakeDevicePool("draft", layer_num=2)
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        node_a = TreeNode()
        node_a.host_len = 16
        node_a.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            draft_host_indices=torch.tensor([300, 301, 302, 303], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )
        node_b = TreeNode()
        node_b.host_len = 16
        node_b.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            draft_host_indices=torch.tensor([400, 401, 402, 403], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        device_indices_a = controller.load_cp([node_a], node_id=201)
        device_indices_b = controller.load_cp([node_b], node_id=202)
        controller.start_loading()

        self.assertEqual(device_indices_a.tolist(), list(range(64, 80)))
        self.assertEqual(device_indices_b.tolist(), list(range(80, 96)))
        self.assertEqual(
            allocator.owner_alloc_calls,
            [[3, 0, 1, 2], [3, 0, 1, 2]],
        )
        self.assertEqual(len(host_pool.loads), 2)
        self.assertEqual([load[2] for load in host_pool.loads], [0, 1])
        for host_indices, device_indices, _, device_pool in host_pool.loads:
            self.assertEqual(
                host_indices.tolist(),
                [100, 101, 102, 103, 200, 201, 202, 203],
            )
            self.assertEqual(device_indices.tolist(), [20, 21, 22, 23, 24, 25, 26, 27])
            self.assertIs(device_pool, allocator.device_pool)

        self.assertEqual(len(draft_host_pool.loads), 2)
        self.assertEqual([load[2] for load in draft_host_pool.loads], [0, 1])
        for host_indices, device_indices, _, device_pool in draft_host_pool.loads:
            self.assertEqual(
                host_indices.tolist(),
                [300, 301, 302, 303, 400, 401, 402, 403],
            )
            self.assertEqual(device_indices.tolist(), [20, 21, 22, 23, 24, 25, 26, 27])
            self.assertIs(device_pool, draft_device_pool)

        self.assertEqual(len(controller.ack_load_queue), 1)
        self.assertEqual(controller.ack_load_queue[0].node_ids, [201, 202])

    def test_cp_start_loading_keeps_zero_owned_load_ack_in_batched_load(self):
        class SequentialOwnerAllocator(FakeAllocator):
            def __init__(self, alloc_results):
                super().__init__()
                self.alloc_results = [result.clone() for result in alloc_results]
                self.device_pool = FakeDevicePool("target", layer_num=2)

            def alloc_pages_with_owners(self, page_owners):
                owners = list(page_owners)
                self.owner_alloc_calls.append(owners)
                if not self.alloc_results:
                    return None
                return self.alloc_results.pop(0).clone()

        host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
        allocator = SequentialOwnerAllocator(
            [
                torch.arange(64, 68, dtype=torch.int64),
                torch.arange(80, 96, dtype=torch.int64),
            ]
        )
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        zero_owned_node = TreeNode()
        zero_owned_node.host_len = 4
        zero_owned_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.tensor([0], dtype=torch.int8),
            page_size=4,
        )
        owned_node = TreeNode()
        owned_node.host_len = 16
        owned_node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        zero_visible_indices = controller.load_cp([zero_owned_node], node_id=301)
        owned_visible_indices = controller.load_cp([owned_node], node_id=302)
        controller.start_loading()

        self.assertEqual(zero_visible_indices.tolist(), [64, 65, 66, 67])
        self.assertEqual(owned_visible_indices.tolist(), list(range(80, 96)))
        self.assertEqual(allocator.owner_alloc_calls, [[0], [3, 0, 1, 2]])
        self.assertEqual(len(host_pool.loads), 2)
        self.assertEqual([load[2] for load in host_pool.loads], [0, 1])
        for host_indices, device_indices, _, _ in host_pool.loads:
            self.assertEqual(host_indices.tolist(), [200, 201, 202, 203])
            self.assertEqual(device_indices.tolist(), [24, 25, 26, 27])
        self.assertEqual(len(controller.ack_load_queue), 1)
        self.assertEqual(controller.ack_load_queue[0].node_ids, [301, 302])

    def test_cp_load_allocates_full_logical_locs_and_transfers_owned_physical_locs(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        device_indices = controller.load_cp([node], node_id=11)
        controller.start_loading()

        self.assertEqual(device_indices.tolist(), list(range(64, 80)))
        self.assertEqual(allocator.alloc_calls, [])
        self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])

    def test_cp_load_returns_valid_locs_while_transferring_padded_tail_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 72, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 6
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=6,
            padded_len=8,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0], dtype=torch.int8),
            page_size=4,
        )

        device_indices = controller.load_cp([node], node_id=112)
        controller.start_loading()

        self.assertEqual(device_indices.tolist(), list(range(64, 70)))
        self.assertEqual(allocator.owner_alloc_calls, [[3, 0]])
        self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])

    def test_start_loading_prepares_load_op_once_for_all_layers(self):
        class PreparingFakeHostPool(FakeHostPool):
            def __init__(self, alloc_result):
                super().__init__(alloc_result)
                self.begin_calls = []
                self.end_calls = 0
                self.active = False

            def begin_load_to_device_op(self, host_indices, device_indices, io_backend):
                self.begin_calls.append(
                    (host_indices.clone(), device_indices.clone(), io_backend)
                )
                self.active = True

            def end_load_to_device_op(self):
                self.end_calls += 1
                self.active = False

            def load_to_device_per_layer(
                self, device_pool, host_indices, device_indices, layer_id, io_backend
            ):
                assert self.active
                super().load_to_device_per_layer(
                    device_pool, host_indices, device_indices, layer_id, io_backend
                )

        host_pool = PreparingFakeHostPool(
            torch.tensor([100, 101, 102, 103], dtype=torch.int64)
        )
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        allocator.device_pool = FakeDevicePool(layer_num=3)
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        controller.load_cp([node], node_id=113)
        controller.start_loading()

        self.assertEqual(len(host_pool.begin_calls), 1)
        self.assertEqual(host_pool.end_calls, 1)
        self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])

    def test_start_loading_emits_descriptor_timing_when_enabled(self):
        host_pool = FakeHostPool(
            torch.tensor([100, 101, 102, 103], dtype=torch.int64)
        )
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        allocator.device_pool = FakeDevicePool(layer_num=2)
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )
        controller.load_cp([node], node_id=114)
        timing_keys = []

        def record_timing(key, start_time, message, *args):
            timing_keys.append(key)

        with patch(
            "sglang.srt.managers.cache_controller.envs."
            "SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get",
            return_value=True,
        ), patch(
            "sglang.srt.managers.cache_controller."
            "_cp_shared_kv_bs_gt1_cache_timing",
            side_effect=record_timing,
        ):
            controller.start_loading()

        self.assertIn("prepare_load_descriptor", timing_keys)
        self.assertIn("submit_h2d_layer_loop", timing_keys)
        self.assertIn("submit_h2d_layer_per_call_slow", timing_keys)
        self.assertIn("end_load_descriptor", timing_keys)

    def test_cp_load_frees_unexpected_owner_allocator_length(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        with self.assertRaisesRegex(
            RuntimeError, "alloc_pages_with_owners returned unexpected length"
        ):
            controller.load_cp([node], node_id=111)

        self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
        self.assertEqual(allocator.frees[0].tolist(), list(range(64, 76)))

    def test_cp_load_with_draft_pool_restores_target_and_draft_locs(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        draft_device_pool = FakeDevicePool("draft")
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        device_indices = controller.load_cp([node], node_id=14)
        controller.start_loading()

        self.assertEqual(device_indices.tolist(), list(range(64, 80)))
        self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
        self.assertEqual(draft_host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
        self.assertIs(draft_host_pool.loads[0][3], draft_device_pool)
        self.assertEqual(len(controller.ack_load_queue), 1)
        self.assertEqual(controller.ack_load_queue[0].node_ids, [14])

    def test_cp_load_valid_tail_with_draft_returns_valid_locs_only(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        draft_device_pool = FakeDevicePool("draft")
        allocator = FakeAllocator(alloc_result=torch.arange(8, 12, dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=draft_device_pool,
        )
        node = TreeNode()
        node.host_len = 3
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=3,
            padded_len=4,
            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            page_owners=torch.tensor([1], dtype=torch.int8),
            page_size=4,
        )

        device_indices = controller.load_cp([node], node_id=178)
        controller.start_loading()

        self.assertEqual(device_indices.tolist(), [8, 9, 10])
        self.assertEqual(allocator.owner_alloc_calls, [[1]])
        self.assertEqual(host_pool.loads[0][1].tolist(), [4, 5, 6, 7])
        self.assertEqual(draft_host_pool.loads[0][1].tolist(), [4, 5, 6, 7])
        self.assertIs(draft_host_pool.loads[0][3], draft_device_pool)
        self.assertEqual(len(controller.ack_load_queue), 1)
        self.assertEqual(controller.ack_load_queue[0].node_ids, [178])

    def test_cp_start_loading_loads_draft_before_target_layer_ready(self):
        order = []
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        draft_host_pool = FakeHostPool(
            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
        )
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=1,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        controller.layer_done_counter = RecordingLayerDoneCounter(order)

        def record_target_load(
            device_pool, host_indices, device_indices, layer_id, io_backend
        ):
            order.append(("target", layer_id))

        def record_draft_load(
            device_pool, host_indices, device_indices, layer_id, io_backend
        ):
            order.append(("draft", layer_id))

        host_pool.load_to_device_per_layer = record_target_load
        draft_host_pool.load_to_device_per_layer = record_draft_load
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        controller.load_cp([node], node_id=114)
        controller.start_loading()

        self.assertEqual(order, [("draft", 0), ("target", 0), ("complete", 0)])

    def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3)
        node = TreeNode()
        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.tensor([0], dtype=torch.int8),
            page_size=4,
        )

        device_indices = controller.load_cp([node], node_id=12)
        controller.start_loading()

        self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
        self.assertEqual(allocator.owner_alloc_calls, [[0]])
        self.assertEqual(host_pool.loads, [])
        self.assertEqual(len(controller.ack_load_queue), 1)

    def test_cp_load_zero_owned_rejects_missing_draft_metadata_when_draft_attached(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            allocator=allocator,
            cp_rank=3,
            draft_host_pool=FakeHostPool(torch.tensor([], dtype=torch.int64)),
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        node = TreeNode()
        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.tensor([0], dtype=torch.int8),
            page_size=4,
        )

        with self.assertRaisesRegex(RuntimeError, "draft KV restore requested"):
            controller.load_cp([node], node_id=33)

        self.assertEqual(allocator.owner_alloc_calls, [[0]])
        self.assertEqual(allocator.frees[0].tolist(), [64, 65, 66, 67])
        self.assertEqual(controller.load_queue, [])
        self.assertEqual(controller.draft_load_queue, [])

    def test_cp_load_queues_cpu_host_indices_before_backend_moves(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        host_indices = HostIndicesTensor(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        node = TreeNode()
        node.host_len = 16
        metadata = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )
        metadata.host_indices = host_indices
        node.cp_hicache = metadata

        controller.load_cp([node], node_id=13)

        queued_op = controller.load_queue[0]
        self.assertEqual(queued_op.host_indices.device.type, "cpu")
        self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])

    def test_direct_layer_page_first_move_indices_keeps_host_order_and_cpu_device_indices(self):
        host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
        host_pool.layout = "layer_page_first"
        controller = self.make_controller(host_pool, cp_rank=1)
        op = CacheOperation(
            host_indices=torch.tensor([12, 8, 9, 10], dtype=torch.int64),
            device_indices=torch.tensor([32, 28, 29, 30], dtype=torch.int64),
            node_id=7,
        )

        host_indices, device_indices = controller.move_indices(op, host_pool)

        self.assertEqual(host_indices.tolist(), [12, 8, 9, 10])
        self.assertEqual(host_indices.device.type, "cpu")
        self.assertEqual(device_indices.tolist(), [32, 28, 29, 30])
        self.assertEqual(device_indices.device.type, "cpu")

    def test_cp_load_rejects_non_contiguous_physical_device_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(
            alloc_result=torch.tensor(
                [64, 65, 66, 67, 68, 69, 71, 70, 72, 73, 74, 75, 76, 77, 78, 79],
                dtype=torch.int64,
            )
        )
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        with self.assertRaisesRegex(
            ValueError, "physical_device_indices.*contiguous page spans"
        ):
            controller.load_cp([node], node_id=31)

    def test_cp_load_rejects_non_contiguous_host_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
            page_size=4,
        )

        with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
            controller.load_cp([node], node_id=32)

    def test_cp_evict_host_frees_target_and_draft_host_indices(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        metadata = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
            page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8),
            page_size=4,
        )

        freed = controller.evict_cp_host(metadata)

        self.assertEqual(freed, 4)
        self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
        self.assertEqual(draft_host_pool.frees[0].tolist(), [200, 201, 202, 203])

    def test_cp_evict_host_rejects_missing_draft_metadata_before_target_free(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(
            host_pool,
            draft_host_pool=draft_host_pool,
            draft_mem_pool_device=FakeDevicePool("draft"),
        )
        metadata = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
            page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8),
            page_size=4,
        )

        with self.assertRaisesRegex(RuntimeError, "draft.*evict.*draft_host_indices"):
            controller.evict_cp_host(metadata)

        self.assertEqual(host_pool.frees, [])
        self.assertEqual(draft_host_pool.frees, [])


if __name__ == "__main__":
    main()