From 1a3bea77f2abadb2206ba93293ee4fc4da76742a Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Sun, 2 Aug 2026 08:35:27 -0700 Subject: [PATCH] [Feat] DCP + HiCache L2 Support (ported from kimi-k3) (#33112) Co-authored-by: Claude Opus 5 (1M context) --- .../scheduler_components/metrics_reporter.py | 7 +- python/sglang/srt/mem_cache/hiradix_cache.py | 5 + .../hybrid_cache/hybrid_pool_assembler.py | 9 + .../sglang/srt/mem_cache/memory_pool_host.py | 3 + python/sglang/srt/mem_cache/pool_host/base.py | 48 +++- python/sglang/srt/mem_cache/pool_host/mla.py | 14 ++ python/sglang/srt/server_args.py | 43 ++++ .../test_unified_radix_cache_kl_dcp.py | 112 ++++++++++ .../mem_cache/test_hicache_dcp_host_pool.py | 210 ++++++++++++++++++ 9 files changed, 442 insertions(+), 9 deletions(-) create mode 100644 test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py create mode 100644 test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 0e5973627..34b7ec9aa 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -1008,10 +1008,9 @@ class SchedulerMetricsReporter: self.scheduler.tree_cache, "token_to_kv_pool_host", None ) or getattr(self.scheduler.tree_cache, "full_kv_pool_host", None) assert host_pool is not None, "Host pool not found" - self.stats.hicache_host_used_tokens = ( - host_pool.size - host_pool.available_size() - ) - self.stats.hicache_host_total_tokens = host_pool.size + host_total = host_pool.logical_size + self.stats.hicache_host_used_tokens = host_total - host_pool.available_size() + self.stats.hicache_host_total_tokens = host_total def _update_lora_metrics(self): """Update LoRA pool metrics for monitoring and autoscaling.""" diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 86811cabf..5ebac56f3 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -99,6 +99,9 @@ class HiRadixCache(RadixCache): # Filled by attach_hybrid_minimax_sparse_pool_to_hiradix_cache. self.token_to_kv_pool_host = None elif isinstance(self.kv_cache, MLATokenToKVPool): + from sglang.srt.runtime_context import get_parallel + + _parallel = get_parallel() self.token_to_kv_pool_host = MLATokenToKVPoolHost( self.kv_cache, server_args.hicache_ratio, @@ -106,6 +109,8 @@ class HiRadixCache(RadixCache): self.page_size, server_args.hicache_mem_layout, allocator_type=allocator_type, + dcp_size=_parallel.attn_dcp_size, + dcp_rank=_parallel.attn_dcp_rank, ) else: raise ValueError("HiRadixCache only supports MHA, MLA, DSA, and MSA models") diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 8fd6d421a..8faec5d77 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -28,6 +28,7 @@ from sglang.srt.mem_cache.pool_host.mha import ( ) from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost from sglang.srt.mem_cache.unified_cache.components import ComponentType +from sglang.srt.runtime_context import get_parallel if TYPE_CHECKING: import torch @@ -72,6 +73,14 @@ def build_kv_host_pool( kwargs = {} if override_kv_cache_dim is not None: kwargs["override_kv_cache_dim"] = override_kv_cache_dim + parallel = get_parallel() + if parallel.dcp_enabled: + assert use_mla, ( + "HiCache + DCP is only wired for the MLA host pool; the MHA host " + "pool has no DCP index translation." + ) + kwargs["dcp_size"] = parallel.attn_dcp_size + kwargs["dcp_rank"] = parallel.attn_dcp_rank return kv_host_pool_cls( kv_pool, server_args.hicache_ratio, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 42c84ece9..041f34a32 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -650,6 +650,8 @@ class LogicalHostPool: f"got size={size}, page_size={page_size}" ) self.size = size + # Stands in for a host pool (and group anchor); DCP never widens it. + self.logical_size = size self.page_size = page_size self.device = "cpu" self.layout = layout @@ -1528,6 +1530,7 @@ class HostPoolGroup: self.page_size = self.anchor_entry.host_pool.page_size self.device = self.anchor_entry.host_pool.device self.size = self.anchor_entry.host_pool.size + self.logical_size = self.anchor_entry.host_pool.logical_size child_write_back_jit = [ getattr(entry.host_pool, "can_use_write_back_jit", False) for entry in entries diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 7018bc579..19770824e 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -79,6 +79,8 @@ def synchronized(func): class HostKVCache(abc.ABC): + dcp_size = 1 + dcp_rank = 0 def __init__( self, @@ -90,9 +92,19 @@ class HostKVCache(abc.ABC): pin_memory: bool, device: str, allocator_type: str = "default", + dcp_size: int = 1, + dcp_rank: int = 0, ): self.device_pool = device_pool - self.page_size = page_size + # page_size arrives widened (x dcp_size); size/page_size/page_num are physical. + self.dcp_size = dcp_size + self.dcp_rank = dcp_rank + assert page_size % dcp_size == 0, ( + f"HiCache host pool page_size ({page_size}) must be a multiple of " + f"dcp_size ({dcp_size}); expected the widened page from the DCP " + "paged allocator." + ) + self.page_size = page_size // dcp_size self.layout = layout self.pin_memory = pin_memory self.device = device @@ -265,16 +277,16 @@ class HostKVCache(abc.ABC): def clear(self): # Initialize memory states and tracking structures. self.mem_state = torch.zeros( - (self.size,), dtype=torch.uint8, device=self.device + (self.logical_size,), dtype=torch.uint8, device=self.device ) - self.free_slots = torch.arange(self.size, dtype=torch.int64) + self.free_slots = torch.arange(self.logical_size, dtype=torch.int64) # Keep freed chunks aside and consume them lazily from alloc() to avoid # concatenating a large free-list on every host-pool free. self.release_slots = [] self.num_release_slots = 0 # Per-slot flag used to detect double-free. # slot_used[k] is true if slot k is allocated. - self.slot_used = torch.zeros(self.size, dtype=torch.bool) + self.slot_used = torch.zeros(self.logical_size, dtype=torch.bool) def available_size(self): return len(self.free_slots) + self.num_release_slots @@ -291,10 +303,36 @@ class HostKVCache(abc.ABC): self.release_slots = [] self.num_release_slots = 0 + @property + def logical_size(self) -> int: + """Slots the radix/controller layer sees: dcp_size of them share a row.""" + return self.size * self.dcp_size + + @property + def logical_page_size(self) -> int: + """Page size in that same logical space (the widened DCP page).""" + return self.page_size * self.dcp_size + + def dcp_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor: + """Transfer kernels index per-rank rows; callers hold widened logical slots. + + Keep this rank's slots (% dcp_size == dcp_rank), then collapse (// dcp_size). + """ + if self.dcp_size == 1: + return indices + owned = indices[indices % self.dcp_size == self.dcp_rank] // self.dcp_size + assert owned.numel() * self.dcp_size == indices.numel(), ( + "HiCache DCP translation expects runs of whole widened pages " + f"(every residue class equally represented); got {indices.numel()} " + f"logical slots -> {owned.numel()} owned rows with dcp_size=" + f"{self.dcp_size}." + ) + return owned + @synchronized def alloc(self, need_size: int) -> Optional[torch.Tensor]: assert ( - need_size % self.page_size == 0 + need_size % self.logical_page_size == 0 ), "The requested size should be a multiple of the page size." if need_size > self.available_size(): return None diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index bbfe420b5..7066b6348 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/python/sglang/srt/mem_cache/pool_host/mla.py @@ -62,6 +62,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): device: str = "cpu", allocator_type: str = "default", override_kv_cache_dim: Optional[int] = None, + dcp_size: int = 1, + dcp_rank: int = 0, ): self.override_kv_cache_dim = override_kv_cache_dim super().__init__( @@ -73,6 +75,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): pin_memory, device, allocator_type, + dcp_size=dcp_size, + dcp_rank=dcp_rank, ) # The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only # helpers in hicache.cuh are guarded by USE_ROCM and the staged @@ -226,6 +230,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): ): if not self._is_device_layer_owned(device_pool, layer_id): return + host_indices = self.dcp_kernel_indices(host_indices) + device_indices = self.dcp_kernel_indices(device_indices) host_layer = self._host_layer_index(layer_id) if io_backend == "kernel": @@ -311,6 +317,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): def _backup_from_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): + # Indices arrive already translated by backup_from_device_all_layer. host_layer = self._host_layer_index(layer_id) if io_backend == "kernel": if self.layout == "layer_first": @@ -370,6 +377,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): + host_indices = self.dcp_kernel_indices(host_indices) + device_indices = self.dcp_kernel_indices(device_indices) if self._is_device_layer_sharded(device_pool): for layer_id in self._owned_device_layer_ids(device_pool): self._backup_from_device_per_layer( @@ -459,6 +468,11 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): raise ValueError(f"Unsupported IO backend: {io_backend}") def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + assert self.dcp_size == 1, ( + "HiCache L3 storage paths are not yet DCP-aware (per-rank shards " + "need dcp_rank-scoped keys); --hicache-storage-backend with " + "--dcp-size > 1 should have been rejected at server start." + ) if self.layout == "layer_first": data_page = self.kv_buffer[:, index : index + self.page_size, :, :] elif self.layout == "page_first": diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 86eedf3fd..facbb2248 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7120,6 +7120,49 @@ class ServerArgs: # Step 2: Storage-layout normalization without changing io backend. self._resolve_storage_layout_compatibility() + # Step 3: DCP compatibility for the L2 (device<->host) path. + self._resolve_hicache_dcp_compatibility() + + def _resolve_hicache_dcp_compatibility(self): + if self.dcp_size <= 1 or not self.enable_hierarchical_cache: + return + if self.hicache_storage_backend is not None: + raise NotImplementedError( + "--hicache-storage-backend (L3) with --dcp-size > 1 is not " + "supported yet: under DCP each rank holds a distinct " + "interleaved MLA KV shard, so the rank-0-only replicated-MLA " + "backup and the storage keys must become dcp_rank-aware " + "first. Run HiCache+DCP with L1/L2 only." + ) + if self.speculative_algorithm is not None: + raise NotImplementedError( + "HiCache with --dcp-size > 1 does not support speculative " + "decoding yet (the draft-model host pool has no DCP index " + "translation)." + ) + if self.enable_lmcache: + raise NotImplementedError( + "--enable-lmcache with --dcp-size > 1 is not supported: " + "LMCache has no DCP-aware index translation." + ) + if self.enable_hisparse: + raise NotImplementedError( + "--enable-hisparse with --dcp-size > 1 is not supported: the " + "HiSparse host pool is constructed without DCP translation." + ) + if not self.use_mla_backend(): + raise NotImplementedError( + "HiCache with --dcp-size > 1 is only supported for MLA models: " + "the index translation lives in MLATokenToKVPoolHost, and the " + "MHA host pool has none." + ) + logger.info( + "HiCache + DCP enabled (L1/L2 only): host pool uses widened " + "logical slot accounting with per-rank physical translation at " + "the transfer boundary (dcp_size=%d).", + self.dcp_size, + ) + def _resolve_layout_io_compatibility(self): if ( self.hicache_mem_layout == "page_first_direct" diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py new file mode 100644 index 000000000..2b6e35a2f --- /dev/null +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py @@ -0,0 +1,112 @@ +"""HiCache L2 under decode context parallelism (DCP) + UnifiedRadixCache. + +Under DCP the radix layer allocates widened logical indices while each rank's +buffers hold only its 1/dcp_size shard, so a missing translation makes cache +hits return another rank's KV. The KL cases catch that as a large divergence. + +Blackwell-only: the MLA DCP decode path needs ``tokenspeed_mla`` (SM100/12x). +""" + +import subprocess +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin +from sglang.test.kl_multiturn_utils import ( + get_input_ids, + make_mamba_decode_assert, + make_mamba_prefill_assert, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=1500, stage="extra-b", runner_config="4-gpu-b200") + +KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" +DCP_SIZE = 4 +PAGE_SIZE = 64 +WIDENED_PAGE = PAGE_SIZE * DCP_SIZE +MAX_MAMBA_CACHE_SIZE = 256 +# Bound the host pools directly. Sizing them off the device pool would need +# --max-total-tokens, which triggers out-of-range KV writes under DCP. +HICACHE_SIZE_GB = 10 + + +class TestUnifiedKimiLinearDcpHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """Kimi Linear + DCP4 + HiCache L2 + UnifiedRadixCache.""" + + kl_threshold = 0.01 + gsm8k_threshold = 0.85 + mmlu_threshold = 0.4 + prefill_cache_assert = staticmethod( + make_mamba_prefill_assert(chunk_size=WIDENED_PAGE) + ) + decode_cache_assert = staticmethod( + make_mamba_decode_assert(track_interval=WIDENED_PAGE) + ) + + @classmethod + def setUpClass(cls): + cls.model = KIMI_LINEAR_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3, + other_args=[ + "--trust-remote-code", + "--tp-size", + "4", + "--dcp-size", + str(DCP_SIZE), + "--page-size", + str(PAGE_SIZE), + "--attention-backend", + "tokenspeed_mla", + "--kv-cache-dtype", + "fp8_e4m3", + "--dcp-comm-backend", + "a2a", + "--dcp-replicate-q-proj", + "--dtype", + "bfloat16", + "--random-seed", + "0", + "--cuda-graph-max-bs-decode", + "64", + "--cuda-graph-backend-prefill", + "disabled", + "--mem-fraction-static", + "0.80", + "--enable-hierarchical-cache", + "--hicache-size", + str(HICACHE_SIZE_GB), + "--hicache-write-policy", + "write_through", + "--max-running-requests", + "64", + "--max-mamba-cache-size", + str(MAX_MAMBA_CACHE_SIZE), + "--enable-metrics", + ], + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18, trust_remote_code=True) + + @classmethod + def tearDownClass(cls): + cls.process.terminate() + try: + cls.process.wait(timeout=60) + except subprocess.TimeoutExpired: + pass + kill_process_tree(cls.process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py b/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py new file mode 100644 index 000000000..66fc2a7b6 --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py @@ -0,0 +1,210 @@ +"""HiCache under decode context parallelism (DCP): host-pool index math. + +Under DCP the radix/controller layer works in a widened logical index space +(page_size * dcp_size wide pages, dcp_size * physical capacity), while each +rank's device and host buffers only materialize the owned 1/dcp_size token +shard (owner rule: index % dcp_size == dcp_rank, physical row = index // +dcp_size — the same rule the device-side KV write and page-table kernels +use). These tests cover the translation helper, the logical/physical host +pool sizing, and that the transfer entry points hand *physical* rows to the +kernels. +""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +DCP_SIZE = 8 +PHYSICAL_PAGE = 64 +WIDENED_PAGE = PHYSICAL_PAGE * DCP_SIZE + + +def _fake_mla_device_pool(size: int = 1024) -> SimpleNamespace: + return SimpleNamespace( + size=size, + store_dtype=torch.float16, + kv_lora_rank=8, + qk_rope_head_dim=4, + layer_num=2, + start_layer=0, + end_layer=1, + device="cpu", + layers_to_capture=None, + layer_shard_enabled=False, + ) + + +def _make_host_pool(dcp_rank: int, device_size: int = 1024) -> MLATokenToKVPoolHost: + return MLATokenToKVPoolHost( + _fake_mla_device_pool(device_size), + host_to_device_ratio=2.0, + host_size=0, + page_size=WIDENED_PAGE, + layout="layer_first", + pin_memory=False, + device="cpu", + dcp_size=DCP_SIZE, + dcp_rank=dcp_rank, + ) + + +class TestDcpKernelIndices(CustomTestCase): + def _bare_pool(self, dcp_size: int, dcp_rank: int) -> MLATokenToKVPoolHost: + pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) + pool.dcp_size = dcp_size + pool.dcp_rank = dcp_rank + return pool + + def test_identity_without_dcp(self): + pool = self._bare_pool(1, 0) + indices = torch.arange(37) + self.assertIs(pool.dcp_kernel_indices(indices), indices) + + def test_aligned_page_translates_to_full_physical_page(self): + # One widened page starting at logical 512 covers physical rows + # 64..127 on every rank. + indices = torch.arange(WIDENED_PAGE, 2 * WIDENED_PAGE) + for rank in range(DCP_SIZE): + pool = self._bare_pool(DCP_SIZE, rank) + out = pool.dcp_kernel_indices(indices) + torch.testing.assert_close( + out, torch.arange(PHYSICAL_PAGE, 2 * PHYSICAL_PAGE) + ) + + def test_matches_owner_rule_on_merged_unordered_pages(self): + # Concatenation of non-adjacent widened pages in arbitrary order, as + # produced by merged CacheOperations after allocator churn. + pages = [3, 0, 5] + indices = torch.cat( + [torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) for p in pages] + ) + for rank in range(DCP_SIZE): + pool = self._bare_pool(DCP_SIZE, rank) + out = pool.dcp_kernel_indices(indices) + expected = ( + indices[indices % DCP_SIZE == rank] // DCP_SIZE + ) # owner rule, same as filter_dcp_local_kv_indices + torch.testing.assert_close(out, expected) + self.assertEqual(out.numel() * DCP_SIZE, indices.numel()) + + def test_ragged_run_is_rejected(self): + pool = self._bare_pool(DCP_SIZE, 0) + with self.assertRaises(AssertionError): + pool.dcp_kernel_indices(torch.arange(WIDENED_PAGE + 1)) + + def test_positional_residue_pairing_survives_host_sort(self): + # move_indices (direct/layer_first) sorts host indices and permutes + # device indices to match. Independent residue filtering of both + # tensors must keep the same token positions on every rank. + g = torch.Generator().manual_seed(0) + host_pages = [7, 2] + device_pages = [1, 4] + host = torch.cat( + [torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) for p in host_pages] + ) + device = torch.cat( + [ + torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) + for p in device_pages + ] + ) + # token identity: position i pairs host[i] <-> device[i] + perm = torch.randperm(host.numel(), generator=g) + # sort host as move_indices does, permuting device alongside + host_sorted, order = host[perm].sort() + device_matched = device[perm][order] + for rank in range(DCP_SIZE): + pool = self._bare_pool(DCP_SIZE, rank) + host_mask = host_sorted % DCP_SIZE == rank + device_mask = device_matched % DCP_SIZE == rank + # same positions selected on both sides -> pairing preserved + torch.testing.assert_close(host_mask, device_mask) + self.assertEqual( + pool.dcp_kernel_indices(host_sorted).numel(), + host.numel() // DCP_SIZE, + ) + + +class TestHostPoolSizingUnderDcp(CustomTestCase): + def test_logical_and_physical_sizing(self): + pool = _make_host_pool(dcp_rank=3) + # kernel-facing page is physical + self.assertEqual(pool.page_size, PHYSICAL_PAGE) + self.assertEqual(pool.logical_page_size, WIDENED_PAGE) + # physical rows = ratio * device physical size, page aligned + self.assertEqual(pool.size, pool.page_num * PHYSICAL_PAGE) + self.assertEqual(pool.logical_size, pool.size * DCP_SIZE) + # buffers materialize physical rows only + self.assertEqual(pool.kv_buffer.shape[1], pool.size) + # allocator surface is logical + self.assertEqual(pool.free_slots.numel(), pool.logical_size) + self.assertEqual(pool.mem_state.numel(), pool.logical_size) + + def test_alloc_is_widened_page_granular(self): + pool = _make_host_pool(dcp_rank=0) + out = pool.alloc(WIDENED_PAGE) + self.assertEqual(out.numel(), WIDENED_PAGE) + with self.assertRaises(AssertionError): + pool.alloc(PHYSICAL_PAGE) # not a multiple of the widened page + + def test_non_dcp_pool_unchanged(self): + pool = MLATokenToKVPoolHost( + _fake_mla_device_pool(), + host_to_device_ratio=2.0, + host_size=0, + page_size=PHYSICAL_PAGE, + layout="layer_first", + pin_memory=False, + device="cpu", + ) + self.assertEqual(pool.page_size, PHYSICAL_PAGE) + self.assertEqual(pool.logical_size, pool.size) + self.assertEqual(pool.logical_page_size, PHYSICAL_PAGE) + + +class TestTransferEntryPointsTranslate(CustomTestCase): + def _run_backup(self, pool, host_indices, device_indices): + device_pool = SimpleNamespace( + data_ptrs=torch.zeros(2, dtype=torch.uint64), + kv_buffer=[torch.zeros(1)] * 2, + ) + # create=True: mla.py imports the kernel only under `if _is_cuda or + # _is_hip`, so the name is absent on the CPU runner this test targets. + with mock.patch( + "sglang.srt.mem_cache.pool_host.mla.transfer_kv_all_layer_mla", + create=True, + ) as kernel: + pool.can_use_jit = False + pool.can_use_write_back_jit = False + with mock.patch.object( + MLATokenToKVPoolHost, "_is_device_layer_sharded", return_value=False + ): + pool.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend="kernel" + ) + return kernel.call_args.kwargs + + def test_backup_receives_physical_rows(self): + pool = _make_host_pool(dcp_rank=5) + logical = torch.arange(2 * WIDENED_PAGE) + kwargs = self._run_backup(pool, logical, logical.clone()) + expected = torch.arange(2 * PHYSICAL_PAGE) + torch.testing.assert_close(kwargs["src_indices"], expected) + torch.testing.assert_close(kwargs["dst_indices"], expected) + + def test_l3_data_page_is_guarded(self): + pool = _make_host_pool(dcp_rank=0) + with self.assertRaises(AssertionError): + pool.get_data_page(0) + + +if __name__ == "__main__": + unittest.main()