diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index fe2916b7f..e229b418e 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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]: diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index d1b65c2db..440d46401 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/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( 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 66fc2a7b6..6cb4b8e18 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 @@ -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, )