[Fix][DCP] Localize widened KV ids in MLA retraction CPU backup/restore (#39487)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

python/sglang/srt/layers/dcp/layout.py

@@ -41,6 +41,19 @@ def get_dcp_lens(
     return torch.clamp((remaining + dcp_size - 1) // dcp_size, min=0)
 
 
+def maybe_dcp_kernel_indices(
+    indices: torch.Tensor, dcp_size: int, dcp_rank: int
+) -> torch.Tensor:
+    """Widened logical slots -> this rank's physical rows.
+
+    Owner rule: slot % dcp_size == dcp_rank, row = slot // dcp_size. The run
+    starts page-aligned, so a strided view selects the owned slots without a mask.
+    """
+    if dcp_size == 1:
+        return indices
+    return indices[dcp_rank::dcp_size] // dcp_size
+
+
 def filter_dcp_local_kv_indices(kv_indices: torch.Tensor):
     """Keep this rank's share of a read-index tensor, still WIDENED.
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -52,6 +52,7 @@ from sglang.srt.configs.mamba_utils import BaseLinearStateParams
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
+from sglang.srt.layers.dcp.layout import maybe_dcp_kernel_indices
 from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
     UnquantizedKVCacheMethod,
 )
@@ -4610,6 +4611,9 @@ class MLATokenToKVPool(KVCache):
             kv_cache[tgt_loc_flat] = kv_cache[src_loc_flat]
 
     def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
+        indices = maybe_dcp_kernel_indices(
+            indices, self._write_loc_dcp_span, get_parallel().attn_dcp_rank
+        )
         current_platform.synchronize()
         kv_cache_cpu = []
         chunk_size = self.cpu_offloading_chunk_size
@@ -4629,6 +4633,9 @@ class MLATokenToKVPool(KVCache):
     def load_cpu_copy(
         self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
     ):
+        indices = maybe_dcp_kernel_indices(
+            indices, self._write_loc_dcp_span, get_parallel().attn_dcp_rank
+        )
         current_platform.synchronize()
         chunk_size = self.cpu_offloading_chunk_size
         for layer_id in range(self.layer_num):

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

@@ -371,19 +371,6 @@ class HostKVCache(abc.ABC):
         """Page size in that same logical space (the widened DCP page)."""
         return self.page_size * self.dcp_size
 
-    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
-        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 indices[self.dcp_rank :: self.dcp_size] // self.dcp_size
-
     @synchronized
     def alloc(self, need_size: int) -> Optional[torch.Tensor]:
         assert need_size % self.logical_page_size == 0, (

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

@@ -19,6 +19,7 @@ from sglang.kernels.ops.kvcache.hicache import (
 from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
 )
+from sglang.srt.layers.dcp.layout import maybe_dcp_kernel_indices
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.mem_cache.pool_host.base import (
     _WRITE_BACK_STAGING_PAGE_CHUNK,
@@ -653,8 +654,12 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         assert not getattr(self, "_is_dummy", False), (
             "load on a dummy (non-src MLA) host pool"
         )
-        host_indices = self.maybe_dcp_kernel_indices(host_indices)
-        device_indices = self.maybe_dcp_kernel_indices(device_indices)
+        host_indices = maybe_dcp_kernel_indices(
+            host_indices, self.dcp_size, self.dcp_rank
+        )
+        device_indices = maybe_dcp_kernel_indices(
+            device_indices, self.dcp_size, self.dcp_rank
+        )
         # 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
@@ -851,8 +856,12 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         assert not getattr(self, "_is_dummy", False), (
             "backup on a dummy (non-src MLA) host pool"
         )
-        host_indices = self.maybe_dcp_kernel_indices(host_indices)
-        device_indices = self.maybe_dcp_kernel_indices(device_indices)
+        host_indices = maybe_dcp_kernel_indices(
+            host_indices, self.dcp_size, self.dcp_rank
+        )
+        device_indices = maybe_dcp_kernel_indices(
+            device_indices, self.dcp_size, self.dcp_rank
+        )
         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(