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 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.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.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 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 FakeDevicePool: device = "cpu" layer_num = 1 def __init__(self, name="target", layer_num=1): self.name = name self.layer_num = layer_num def register_layer_transfer_counter(self, counter): self.counter = counter 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[0][1].tolist(), [4, 5, 6, 7]) def test_cp_write_rejects_incomplete_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) with self.assertRaisesRegex( ValueError, "_write_cp expects page-aligned device_indices" ): controller.write(logical_locs, node_id=21) 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[0][1].tolist(), [4, 5, 6, 7]) self.assertEqual(draft_host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) self.assertIs(draft_host_pool.backups[0][2], 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_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_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_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_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_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_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()