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

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

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