[Feat] DCP + HiCache L2 Support (ported from kimi-k3) (#33112)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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_total = host_pool.logical_size
-            host_pool.size - host_pool.available_size()
+        self.stats.hicache_host_used_tokens = host_total - host_pool.available_size()
-        )
+        self.stats.hicache_host_total_tokens = host_total
-        self.stats.hicache_host_total_tokens = host_pool.size
 
     def _update_lora_metrics(self):
         """Update LoRA pool metrics for monitoring and autoscaling."""

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")

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,

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

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

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":

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"

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()

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()