From 2c37b90ad6e9c79b7c92be6751cc9d481d6abfe9 Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Mon, 14 Sep 2026 22:33:41 -0700 Subject: [PATCH] [Fix][DCP] Localize widened KV ids in MLA retraction CPU backup/restore (#39487) Co-authored-by: Claude Fable 5.1 --- python/sglang/srt/layers/dcp/layout.py | 13 ++++ python/sglang/srt/mem_cache/memory_pool.py | 7 ++ python/sglang/srt/mem_cache/pool_host/base.py | 13 ---- python/sglang/srt/mem_cache/pool_host/mla.py | 17 +++-- .../test_decode_retraction_backup.py | 69 ++++++++++++++++++- .../mem_cache/test_hicache_dcp_host_pool.py | 24 ++----- 6 files changed, 106 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/layers/dcp/layout.py b/python/sglang/srt/layers/dcp/layout.py index 6fdf99ee9..58697f443 100644 --- a/python/sglang/srt/layers/dcp/layout.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index db9d8654f..28333dffb 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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): diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index fa06cbf5b..7be9f25c3 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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, ( diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 5a1f32726..b63ca9723 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/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( diff --git a/test/registered/unit/mem_cache/test_decode_retraction_backup.py b/test/registered/unit/mem_cache/test_decode_retraction_backup.py index 4b9f456c2..2d0589438 100644 --- a/test/registered/unit/mem_cache/test_decode_retraction_backup.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py b/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py index 29153e139..1fb9254ce 100644 --- a/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py +++ b/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, )