[DCP] Drop the prefill index-selection syncs by taking each rank's rows by stride (#35084)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

@@ -51,6 +51,25 @@ def filter_dcp_local_kv_indices(kv_indices: torch.Tensor):
     return kv_indices
 
 
+def filter_dcp_local_chunk_kv_indices(
+    kv_indices: torch.Tensor,
+    chunk_starts_cpu: torch.Tensor,
+    chunk_seq_lens_cpu: torch.Tensor,
+) -> torch.Tensor:
+    parallel = get_parallel()
+    if not parallel.dcp_enabled:
+        return kv_indices
+
+    dcp_size = parallel.dcp_size
+    parts = []
+    offset = 0
+    for start, length in zip(chunk_starts_cpu.tolist(), chunk_seq_lens_cpu.tolist()):
+        first = (parallel.dcp_rank - start) % dcp_size
+        parts.append(kv_indices[offset + first : offset + length : dcp_size])
+        offset += length
+    return torch.cat(parts) // dcp_size
+
+
 def update_local_kv_lens_for_dcp(kv_len_arr):
     """In-place per-rank KV length: the start=0 case of get_dcp_lens.
 

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

@@ -109,10 +109,9 @@ def prepare_decode_context_parallel_metadata(
         extend_prefix_lens_sum,
         parallel.dcp_size,
     )
+    # Prefix lengths are dcp_size-aligned (widened allocator page), so no nonzero().
     dcp_local_prefix_kv_indices = (
-        dcp_prefix_kv_indices[
+        dcp_prefix_kv_indices[parallel.dcp_rank :: parallel.dcp_size]
-            dcp_prefix_kv_indices % parallel.dcp_size == parallel.dcp_rank
-        ]
         // parallel.dcp_size
     )
     dcp_kv_buffer = torch.empty(

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -10,6 +10,7 @@ from sglang.kernels.ops.kvcache.kv_indices import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.environ import envs
+from sglang.srt.layers.dcp.layout import filter_dcp_local_chunk_kv_indices
 from sglang.srt.model_executor.forward_context import (
     get_req_to_token_pool,
     get_token_to_kv_pool,
@@ -84,6 +85,11 @@ class ForwardBatchDeepSeekMHAMixin:
                 chunk_kv_indices,
                 req_to_token.shape[1],
             )
+            chunk_kv_indices = filter_dcp_local_chunk_kv_indices(
+                chunk_kv_indices,
+                self.prefix_chunk_starts_cpu[idx],
+                self.prefix_chunk_seq_lens_cpu[idx],
+            )
             self.prefix_chunk_kv_indices.append(chunk_kv_indices)
 
     # Here we suppose the length of each chunk is equal

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -450,7 +450,6 @@ class DeepseekMHAForwardMixin:
         forward_batch: ForwardBatch,
     ):
         if _is_cuda:
-            kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices)
             kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer(
                 self.attn_mha, kv_indices, dst_dtype
             )

test/registered/dcp/test_dcp_layout_unit.py

@@ -22,7 +22,10 @@ import torch
 from sglang.srt import runtime_context as rc
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
-from sglang.srt.layers.dcp.layout import get_dcp_lens
+from sglang.srt.layers.dcp.layout import (
+    filter_dcp_local_chunk_kv_indices,
+    get_dcp_lens,
+)
 from sglang.srt.layers.linear import QKVParallelLinear
 from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
 from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
@@ -47,6 +50,97 @@ def _legacy_inplace_formula(length: int, n: int, rank: int) -> int:
     return (length - rank - 1) // n + 1
 
 
+class TestFilterDcpLocalChunkKvIndices(CustomTestCase):
+    PAGE = 64
+
+    def _build_chunk(self, starts, lens, dcp_size, seed=0):
+        g = torch.Generator().manual_seed(seed)
+        widened = self.PAGE * dcp_size
+        runs = []
+        for start, length in zip(starts, lens):
+            if length == 0:
+                runs.append(torch.empty(0, dtype=torch.int64))
+                continue
+            pos = torch.arange(start, start + length)
+            page_of = pos // widened
+            bases = (
+                torch.randint(0, 64, (int(page_of.max()) + 1,), generator=g) * widened
+            )
+            runs.append(bases[page_of] + pos % widened)
+        return torch.cat(runs) if runs else torch.empty(0, dtype=torch.int64)
+
+    def _owner_rule(self, kv, dcp_size, dcp_rank):
+        return kv[kv % dcp_size == dcp_rank] // dcp_size
+
+    def _run(self, starts, lens, dcp_size, dcp_rank, seed=0):
+        kv = self._build_chunk(starts, lens, dcp_size, seed)
+        with rc.get_parallel().override(
+            dcp_enabled=dcp_size > 1, dcp_size=dcp_size, dcp_rank=dcp_rank
+        ):
+            got = filter_dcp_local_chunk_kv_indices(
+                kv, torch.tensor(starts), torch.tensor(lens)
+            )
+        return kv, got
+
+    def test_matches_owner_rule_on_unaligned_runs(self):
+        cases = [
+            ([0], [1]),
+            ([0, 0], [2048, 2048]),
+            ([8192], [3]),
+            ([5, 13, 31], [7, 0, 19]),
+            ([1365, 1365, 1365], [1365, 1365, 1365]),
+            ([43690, 43690], [43690, 17]),
+            ([7, 7, 7, 7, 7], [11, 13, 0, 1, 40]),
+        ]
+        for dcp_size in [2, 3, 4, 8]:
+            for dcp_rank in range(dcp_size):
+                for i, (starts, lens) in enumerate(cases):
+                    kv, got = self._run(starts, lens, dcp_size, dcp_rank, seed=i)
+                    torch.testing.assert_close(
+                        got,
+                        self._owner_rule(kv, dcp_size, dcp_rank),
+                        msg=f"size={dcp_size} rank={dcp_rank} "
+                        f"starts={starts} lens={lens}",
+                    )
+
+    def test_row_count_matches_get_dcp_lens(self):
+        starts, lens = [1365, 2730, 0], [1365, 900, 37]
+        for dcp_size in [2, 3, 4, 8]:
+            for dcp_rank in range(dcp_size):
+                _, got = self._run(starts, lens, dcp_size, dcp_rank)
+                expected = get_dcp_lens(
+                    torch.tensor(lens), dcp_size, dcp_rank, start=torch.tensor(starts)
+                )
+                self.assertEqual(
+                    got.numel(),
+                    int(expected.sum()),
+                    f"size={dcp_size} rank={dcp_rank}",
+                )
+
+    def test_second_chunk_at_batch_three(self):
+        starts, lens = [2730, 2730, 2730], [1365, 1365, 1365]
+        for dcp_rank in range(8):
+            kv, got = self._run(starts, lens, 8, dcp_rank)
+            torch.testing.assert_close(got, self._owner_rule(kv, 8, dcp_rank))
+            naive = kv[dcp_rank::8] // 8
+            self.assertNotEqual(
+                got.numel(),
+                naive.numel(),
+                f"rank={dcp_rank}: phase-free stride happens to match here, "
+                f"so this case no longer guards the phase term",
+            )
+
+    def test_identity_without_dcp(self):
+        kv = torch.arange(37)
+        with rc.get_parallel().override(dcp_enabled=False, dcp_size=1, dcp_rank=0):
+            self.assertIs(
+                filter_dcp_local_chunk_kv_indices(
+                    kv, torch.tensor([0]), torch.tensor([37])
+                ),
+                kv,
+            )
+
+
 class TestGetDcpLens(CustomTestCase):
     def test_start_none_matches_owner_count(self):
         for n in DCP_SIZES: