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

test/registered/unit/mem_cache/test_decode_retraction_backup.py

@@ -9,9 +9,14 @@ from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.common import retraction_backup
 from sglang.srt.mem_cache.hicache_storage import PoolName
 from sglang.srt.mem_cache.kv_cache_builder import maybe_register_hicache_draft
-from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    ReqToTokenPool,
+)
 from sglang.srt.mem_cache.unified_cache.components import ComponentType
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.speculative.base_spec_worker import (
     HiCacheDraftMode,
@@ -215,5 +220,67 @@ class TestDecodeRetractionBackup(unittest.TestCase):
         req_to_token_pool.free(req)
 
 
+DCP_SIZE = 4
+DCP_RANK = 1
+DCP_ROWS = 8
+
+
+def _bare_mla_pool() -> MLATokenToKVPool:
+    pool = object.__new__(MLATokenToKVPool)
+    pool.layer_num = 2
+    pool.cpu_offloading_chunk_size = 3
+    pool.kv_buffer = [
+        (torch.arange(DCP_ROWS, dtype=torch.float32) + 100 * layer).view(DCP_ROWS, 1, 1)
+        for layer in range(pool.layer_num)
+    ]
+    return pool
+
+
+def _dcp():
+    return get_parallel().override(
+        dcp_enabled=True, attn_dcp_size=DCP_SIZE, attn_dcp_rank=DCP_RANK
+    )
+
+
+class TestDcpRetractionBackup(unittest.TestCase):
+    """`req_to_token` names KV slots in the widened DCP id space while
+    `kv_buffer` holds only this rank's rows; a widened id used as a row index
+    reads past the buffer or copies another token's row."""
+
+    def test_restore_lands_on_new_owned_rows(self):
+        pool = _bare_mla_pool()
+        before = [buf.clone() for buf in pool.kv_buffer]
+        old_widened = torch.arange(0, 12, dtype=torch.int64)
+        new_widened = torch.arange(12, 24, dtype=torch.int64)
+
+        with _dcp():
+            pool.load_cpu_copy(pool.get_cpu_copy(old_widened), new_widened)
+
+        old_rows = old_widened[DCP_RANK::DCP_SIZE] // DCP_SIZE
+        new_rows = new_widened[DCP_RANK::DCP_SIZE] // DCP_SIZE
+        self.assertEqual(new_rows.tolist(), [3, 4, 5])
+        untouched = torch.tensor(
+            [r for r in range(DCP_ROWS) if r not in new_rows.tolist()]
+        )
+        for layer in range(pool.layer_num):
+            torch.testing.assert_close(
+                pool.kv_buffer[layer][new_rows], before[layer][old_rows]
+            )
+            torch.testing.assert_close(
+                pool.kv_buffer[layer][untouched], before[layer][untouched]
+            )
+
+    def test_resolved_pool_takes_ids_as_rows(self):
+        pool = _bare_mla_pool()
+        pool.write_loc_is_dcp_resolved = True
+        rows = torch.arange(DCP_ROWS, dtype=torch.int64)
+
+        with _dcp():
+            kv_cpu = pool.get_cpu_copy(rows)
+
+        for layer in range(pool.layer_num):
+            torch.testing.assert_close(torch.cat(kv_cpu[layer]), pool.kv_buffer[layer])
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py

@@ -16,6 +16,7 @@ from unittest import mock
 
 import torch
 
+from sglang.srt.layers.dcp.layout import maybe_dcp_kernel_indices
 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
@@ -57,24 +58,16 @@ def _make_host_pool(dcp_rank: int, device_size: int = 1024) -> MLATokenToKVPoolH
 
 
 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.maybe_dcp_kernel_indices(indices), indices)
+        self.assertIs(maybe_dcp_kernel_indices(indices, 1, 0), 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.maybe_dcp_kernel_indices(indices)
+            out = maybe_dcp_kernel_indices(indices, DCP_SIZE, rank)
             torch.testing.assert_close(
                 out, torch.arange(PHYSICAL_PAGE, 2 * PHYSICAL_PAGE)
             )
@@ -87,19 +80,13 @@ class TestDcpKernelIndices(CustomTestCase):
             [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.maybe_dcp_kernel_indices(indices)
+            out = maybe_dcp_kernel_indices(indices, DCP_SIZE, rank)
             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.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
         # device indices to match. Independent residue filtering of both
@@ -122,13 +109,12 @@ class TestDcpKernelIndices(CustomTestCase):
         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.maybe_dcp_kernel_indices(host_sorted).numel(),
+                maybe_dcp_kernel_indices(host_sorted, DCP_SIZE, rank).numel(),
                 host.numel() // DCP_SIZE,
             )