[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_size == parallel.dcp_rank
-        ]
+        dcp_prefix_kv_indices[parallel.dcp_rank :: parallel.dcp_size]
         // 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
             )