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.memory_pool_host import ( MHATokenToKVPoolHost, MLATokenToKVPoolHost, NSATokenToKVPoolHost, ) 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 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_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) 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_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_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_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_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_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()