[DCP]Localize HiCache DCP indices once per transfer, not per layer (#34889)

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

python/sglang/srt/mem_cache/pool_host/base.py

@@ -340,21 +340,18 @@ class HostKVCache(abc.ABC):
         """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:
+    def maybe_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}."
+        assert indices.numel() % self.dcp_size == 0, (
+            "HiCache DCP translation expects runs of whole widened pages; got "
+            f"{indices.numel()} logical slots with dcp_size={self.dcp_size}."
         )
-        return owned
+        return indices[self.dcp_rank :: self.dcp_size] // self.dcp_size
 
     @synchronized
     def alloc(self, need_size: int) -> Optional[torch.Tensor]:

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -252,8 +252,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
     ):
         if not is_draft and 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_indices = self.maybe_dcp_kernel_indices(host_indices)
+        device_indices = self.maybe_dcp_kernel_indices(device_indices)
         # MTP draft layers do not participate in CP layer sharding.
         host_layer_id = layer_id if is_draft else self._host_layer_index(layer_id)
         device_layer_id = 0 if is_draft else layer_id
@@ -416,8 +416,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)
+        host_indices = self.maybe_dcp_kernel_indices(host_indices)
+        device_indices = self.maybe_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(

test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py

@@ -66,7 +66,7 @@ class TestDcpKernelIndices(CustomTestCase):
     def test_identity_without_dcp(self):
         pool = self._bare_pool(1, 0)
         indices = torch.arange(37)
-        self.assertIs(pool.dcp_kernel_indices(indices), indices)
+        self.assertIs(pool.maybe_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
@@ -74,7 +74,7 @@ class TestDcpKernelIndices(CustomTestCase):
         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)
+            out = pool.maybe_dcp_kernel_indices(indices)
             torch.testing.assert_close(
                 out, torch.arange(PHYSICAL_PAGE, 2 * PHYSICAL_PAGE)
             )
@@ -88,7 +88,7 @@ class TestDcpKernelIndices(CustomTestCase):
         )
         for rank in range(DCP_SIZE):
             pool = self._bare_pool(DCP_SIZE, rank)
-            out = pool.dcp_kernel_indices(indices)
+            out = pool.maybe_dcp_kernel_indices(indices)
             expected = (
                 indices[indices % DCP_SIZE == rank] // DCP_SIZE
             )  # owner rule, same as filter_dcp_local_kv_indices
@@ -98,7 +98,7 @@ class TestDcpKernelIndices(CustomTestCase):
     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))
+            pool.maybe_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
@@ -128,7 +128,7 @@ class TestDcpKernelIndices(CustomTestCase):
             # 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(),
+                pool.maybe_dcp_kernel_indices(host_sorted).numel(),
                 host.numel() // DCP_SIZE,
             )