launcher/dcp2/diffs/09-models_deepseek_v4_1_nvidia_flashinfer_sparse.py.diff

--- baked-ours/models/deepseek_v4_1/nvidia/flashinfer_sparse.py	2026-09-10 20:13:20.000000000 +0200
+++ baked/models/deepseek_v4_1/nvidia/flashinfer_sparse.py	2026-09-11 10:13:47.000000000 +0200
@@ -2,11 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """DeepSeek V4 FlashInfer sparse MLA backend."""
 
+import os
 from typing import TYPE_CHECKING, ClassVar, cast
 
 import torch
 
 from vllm.config import VllmConfig
+from vllm.distributed import get_dcp_group
+from vllm.logger import init_logger
 from vllm.config.cache import CacheDType
 from vllm.forward_context import get_forward_context
 from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
@@ -32,10 +35,94 @@
     get_dspark_swa_index_width,
 )
 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_filter_and_convert_dcp_index,
+)
+from vllm.v1.attention.ops.dcp import (
+    CPTritonContext,
+    cp_lse_ag_out_rs,
+    dcp_a2a_lse_reduce,
+)
 
 if TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
 
+logger = init_logger(__name__)
+
+from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (  # noqa: E402
+    _PREFILL_TIMING as _DSV41_PREFILL_TIMING,
+    _PrefillTimer as _DSV41PrefillTimer,
+    prefill_acc_add as _dsv41_prefill_acc_add,
+)
+
+# Stage 9: DCP prefill attention via compressed-KV gather (no q gather / combine).
+_DCP_PREFILL_GATHER = os.environ.get("DSV41_DCP_PREFILL_GATHER", "0") == "1"
+
+# FlashInfer's SM120 sparse-MLA kernels write this (not -inf) into the LSE of a
+# row with no candidate and no sink; the DCP combine needs -inf there.
+_SM120_EMPTY_LSE = -1e29
+
+
+def _sm120_sparse_attention_with_lse(
+    query: torch.Tensor,
+    swa_kv_cache: torch.Tensor,
+    workspace: torch.Tensor,
+    swa_indices: torch.Tensor,
+    swa_lens: torch.Tensor | None,
+    compressed_kv_cache: torch.Tensor | None,
+    extra_indices: torch.Tensor | None,
+    extra_lens: torch.Tensor | None,
+    out: torch.Tensor,
+    sm_scale: float,
+    sinks: torch.Tensor | None,
+    lse: torch.Tensor,
+) -> None:
+    """FlashInfer's SM120 packed sparse-MLA attention, returning the LSE.
+
+    Mirrors the SM120 branch of ``trtllm_batch_decode_sparse_mla_dsv4``
+    (flashinfer 0.7.0rc1, ``flashinfer/mla/_core.py``), whose public entry
+    hard-codes ``return_lse=False``. ``query``/``out`` are ``[T, H, 512]``
+    bf16, ``lse`` is ``[T, H]`` fp32 and receives the LSE in base 2 of the
+    ``sm_scale``-scaled logits with the sink folded in (``log2(2^lse +
+    2^(sink*log2e))``); rows with no candidate and no sink hold ``-1e30``.
+    """
+    from flashinfer.mla._core import (
+        _check_sm120_dsv4_kv_cache_layout,
+        _SparseMLASegment,
+        _trtllm_batch_decode_sparse_mla_sm120,
+    )
+
+    num_tokens = query.shape[0]
+    swa_kv_cache = _check_sm120_dsv4_kv_cache_layout(swa_kv_cache, "NHD", "swa_kv_cache")
+    segments = [
+        _SparseMLASegment(indices=swa_indices.reshape(num_tokens, -1), lengths=swa_lens)
+    ]
+    if extra_indices is not None:
+        assert compressed_kv_cache is not None
+        compressed_kv_cache = _check_sm120_dsv4_kv_cache_layout(
+            compressed_kv_cache, "NHD", "compressed_kv_cache"
+        )
+        segments.append(
+            _SparseMLASegment(
+                indices=extra_indices.reshape(num_tokens, -1),
+                lengths=extra_lens,
+                kv_cache=compressed_kv_cache,
+            )
+        )
+    _trtllm_batch_decode_sparse_mla_sm120(
+        query=query.unsqueeze(1),
+        kv_cache=swa_kv_cache,
+        workspace_buffer=workspace,
+        sparse_mla_segments=segments,
+        out=out.unsqueeze(1),
+        sm_scale=float(sm_scale),
+        sinks=sinks,
+        lse=lse,
+        return_lse=True,
+        kv_scale_format="auto",
+        kv_cache_format="fp8",
+    )
+
 _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
 
@@ -192,6 +279,23 @@
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # DCP port staging (DCP2-PLAN.md): the sharded-KV attention path
+        # (per-rank sparse attention + LSE combine) lands in stage 3. Until
+        # DSV41_DCP_STAGE >= 3, a DCP>1 boot must fail here -- after the KV
+        # cache config is computed and logged, before any forward -- rather
+        # than run the dcp=1 kernels over a sharded cache and serve garbage.
+        import os
+
+        dcp = self.vllm_config.parallel_config.decode_context_parallel_size
+        if dcp > 1 and int(os.environ.get("DSV41_DCP_STAGE", "3")) < 3:
+            raise NotImplementedError(
+                f"DeepSeek-V4.1 sparse MLA under DCP{dcp}: the attention combine "
+                "(port stage 3) is disabled by DSV41_DCP_STAGE; unset it or run "
+                "with --decode-context-parallel-size 1."
+            )
+
 
 class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
     """SWA metadata for the FlashInfer sparse decode path (varlen decode)."""
@@ -580,6 +684,11 @@
     backend_cls = DeepseekV4FlashInferMLASparseBackend
     swa_backend_cls = DeepseekSparseSWAFlashInferSM120Backend
     use_fp8_ds_mla_layout: ClassVar[bool] = True
+    # DCP: per-rank partial attention with the LSE + cross-rank combine.
+    can_return_lse_for_decode: ClassVar[bool] = True
+    # Under DCP the per-token q gather / combine has no cross-token dependency:
+    # bound the transient [T, 2*heads, 512] buffers inside a prefill chunk.
+    DCP_PREFILL_TOKEN_CHUNK: ClassVar[int] = 2048
 
     @staticmethod
     def _get_workspace(device: torch.device) -> torch.Tensor:
@@ -618,14 +727,56 @@
         from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config
 
         required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size)
-        if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk):
+        # DCP: the kernel sees the gathered heads [rank0 heads | rank1 heads |
+        # ...] in TP order; the local slice must be unpadded so the head
+        # reduce-scatter hands each rank exactly its own heads back.
+        self.n_dcp_heads = self.n_local_heads * self.dcp_world_size
+        if self.dcp_world_size > 1:
+            if self.padded_heads != self.n_local_heads:
+                raise NotImplementedError(
+                    "DeepSeek-V4.1 DCP needs an unpadded local head count in "
+                    f"{_SPARSE_MLA_SUPPORTED_Q_HEADS}; got {self.n_local_heads} "
+                    f"(padded to {self.padded_heads})."
+                )
+            if _pad_to_supported_q_heads(self.n_dcp_heads) != self.n_dcp_heads:
+                raise NotImplementedError(
+                    f"DeepSeek-V4.1 DCP gathered head count {self.n_dcp_heads} is "
+                    f"not a supported sparse-MLA h_q {_SPARSE_MLA_SUPPORTED_Q_HEADS}."
+                )
+        kernel_heads = self.n_dcp_heads if self.dcp_world_size > 1 else self.padded_heads
+        if not has_flashinfer_sparse_mla_sm120_config(kernel_heads, required_topk):
             raise RuntimeError(
                 "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
                 "DSV4 sparse MLA decode specialization for "
-                f"(num_q_heads={self.padded_heads}, top_k={required_topk}). "
+                f"(num_q_heads={kernel_heads}, top_k={required_topk}). "
                 "Install a FlashInfer build containing "
                 "flashinfer-ai/flashinfer#4380."
             )
+        # DCP state: the group (consecutive TP ranks), the gathered sink
+        # (built lazily on the first forward, after weight loading) and the
+        # neutral SWA index rows used on ranks other than 0.
+        self.dcp_group = get_dcp_group() if self.dcp_world_size > 1 else None
+        self._dcp_sink: torch.Tensor | None = None
+        self._dcp_empty_swa: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+        # One Triton launch context per layer for the LSE-combine kernel (a
+        # fresh context per call re-specializes the launch on every step).
+        self._cp_ctx = CPTritonContext() if self.dcp_world_size > 1 else None
+        # Combine flavour: "ag_rs" (LSE all-gather + head reduce-scatter, two
+        # collectives) or "a2a" (one all-to-all of packed output+LSE).
+        self._dcp_combine_a2a = os.environ.get("DSV41_DCP_COMBINE", "ag_rs") == "a2a"
+        # Optional per-phase timing (CUDA events, synchronises; diagnostics
+        # only): DSV41_DCP_TIMING=1 logs gather / kernel / combine ms per call,
+        # averaged over the last 100 calls, on DCP rank 0.
+        self._dcp_timing = os.environ.get("DSV41_DCP_TIMING", "0") == "1"
+        self._dcp_t_acc = [0.0, 0.0, 0.0, 0]
+        if self.dcp_world_size > 1 and self.compress_ratio > 0:
+            logger.info_once(
+                "DeepSeek-V4.1 SM120 attention under DCP%d: %d local heads -> %d "
+                "gathered, SWA window + sink on DCP rank 0, base-2 LSE combine",
+                self.dcp_world_size,
+                self.n_local_heads,
+                self.n_dcp_heads,
+            )
         self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
             self._o_proj_block_size
         )
@@ -762,6 +913,262 @@
             q = padded_query
         return q.contiguous()
 
+    # ---- DCP helpers -------------------------------------------------------
+
+    def _compressed_topk_to_slots(
+        self,
+        topk_indices: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        is_valid_token: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Indexer top-k (compressed ids) -> physical slots + per-row count.
+
+        DCP1: the request-local ids map through the block table as today.
+        DCP>1: the ids are GLOBAL compressed ids (indexer merge / short-context
+        bypass); de-interleave them in state coordinates, drop the states other
+        ranks own and compact this rank's slots to a prefix (-1 tail, count).
+        ``block_size`` is the number of states per page on this rank.
+        """
+        num_tokens = topk_indices.shape[0]
+        if self.dcp_world_size == 1:
+            return compute_global_topk_indices_and_lens(
+                topk_indices,
+                token_to_req_indices,
+                block_table,
+                block_size,
+                is_valid_token,
+            )
+        slots, lens = triton_filter_and_convert_dcp_index(
+            token_to_req_indices[:num_tokens].contiguous(),
+            block_table,
+            topk_indices,
+            dcp_size=self.dcp_world_size,
+            dcp_rank=self.dcp_rank,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            BLOCK_SIZE=block_size,
+            BLOCK_STRIDE_ROWS=block_size,
+            NUM_TOPK_TOKENS=topk_indices.shape[1],
+            BLOCK_N=128,
+            return_valid_counts=True,
+        )
+        # CUDA-graph padding rows attend nothing (parity with the DCP1 kernel).
+        lens.masked_fill_(~is_valid_token[:num_tokens], 0)
+        return slots, lens
+
+    # Stage 9: one prefill step's gathered compressed KV per source cache and
+    # remapped top-k per index buffer, shared by every layer of that step. The
+    # metadata object is held so the cache resets exactly once per step.
+    _dcp_gather_cache: ClassVar[dict] = {"meta": None, "kv": {}, "idx": {}, "geom": None}
+
+    def _dcp_prefill_gather(
+        self,
+        compressed_k_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        token_to_req: torch.Tensor,
+        is_valid_token: torch.Tensor,
+        attn_metadata: DeepseekV4FlashMLAMetadata,
+        num_decodes: int,
+        num_prefills: int,
+        block_size: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """DCP prefill without q gather / LSE combine: gather the DCP ranks'
+        compressed pages of every prefill request into one rank-major paged
+        buffer and remap the GLOBAL top-k ids into it (see stage 9 notes).
+        Returns (gathered cache [W*P_total, block_size, 1, B], slots, lens)."""
+        cache = type(self)._dcp_gather_cache
+        if cache["meta"] is not attn_metadata:
+            cache["meta"] = attn_metadata
+            cache["kv"] = {}
+            cache["idx"] = {}
+            cache["geom"] = None
+        W = self.dcp_world_size
+        K = self.cp_kv_cache_interleave_size
+        bs = block_size
+        dev = topk_indices.device
+        if cache["geom"] is None:
+            assert attn_metadata.seq_lens is not None
+            seq_lens = attn_metadata.seq_lens[num_decodes : num_decodes + num_prefills].tolist()
+            # states that exist after this step (state = pos // ratio) -> virtual
+            # blocks of bs*W states = one page per rank; at least one page so
+            # the collective never carries zero bytes.
+            pages = [max(1, -(-(int(sl) // self.compress_ratio) // (bs * W))) for sl in seq_lens]
+            base = [0]
+            for p in pages:
+                base.append(base[-1] + W * p)
+            cache["geom"] = (
+                pages,
+                torch.tensor(pages, device=dev, dtype=torch.int64),
+                torch.tensor(base[:-1], device=dev, dtype=torch.int64),
+            )
+        pages, pages_t, base_t = cache["geom"]
+        key = compressed_k_cache.data_ptr()
+        kv = cache["kv"].get(key)
+        if kv is None:
+            with _DSV41PrefillTimer("attn_gather", int(topk_indices.shape[0])):
+                pool = self._as_sparse_cache(compressed_k_cache)
+                parts = []
+                for i, p in enumerate(pages):
+                    page_ids = attn_metadata.block_table[num_decodes + i, :p].long()
+                    local = pool.index_select(0, page_ids).contiguous()
+                    parts.append(self.dcp_group.all_gather(local, dim=0))
+                kv = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
+            cache["kv"][key] = kv
+        ikey = (topk_indices.data_ptr(), tuple(topk_indices.shape))
+        idx = cache["idx"].get(ikey)
+        if idx is None:
+            n_req = max(len(pages) - 1, 0)
+            req = (token_to_req.long() - num_decodes).clamp_(0, n_req)
+            ids = topk_indices.long()
+            valid = ids >= 0
+            st = ids.clamp_min(0)
+            span = bs * W
+            vblock = st // span
+            off = st % span
+            owner = (off // K) % W
+            loff = (off // (K * W)) * K + off % K
+            slot = (base_t[req][:, None] + owner * pages_t[req][:, None] + vblock) * bs + loff
+            slots = torch.where(valid, slot, torch.full_like(slot, -1)).to(torch.int32)
+            lens = valid.sum(dim=-1, dtype=torch.int32)
+            lens.masked_fill_(~is_valid_token[: lens.shape[0]], 0)
+            idx = (slots, lens)
+            cache["idx"][ikey] = idx
+        return kv, idx[0], idx[1]
+
+    def _dcp_attn_sink(self) -> torch.Tensor:
+        """The sink for the gathered heads, in TP head order (collective)."""
+        if self._dcp_sink is None:
+            assert self.dcp_group is not None
+            local = self.attn_sink.data[: self.n_local_heads].contiguous()
+            self._dcp_sink = self.dcp_group.all_gather(local, dim=0)
+        return self._dcp_sink
+
+    def _dcp_neutral_swa(
+        self, num_tokens: int, width: int, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """All -1 SWA indices and zero lengths (ranks other than 0)."""
+        buf = self._dcp_empty_swa.get(width)
+        if buf is None or buf[0].shape[0] < num_tokens:
+            rows = max(num_tokens, self.max_num_batched_tokens)
+            buf = (
+                torch.full((rows, 1, width), -1, dtype=torch.int32, device=device),
+                torch.zeros(rows, dtype=torch.int32, device=device),
+            )
+            self._dcp_empty_swa[width] = buf
+        return buf[0][:num_tokens], buf[1][:num_tokens]
+
+    def _run_sm120_dcp(
+        self,
+        q: torch.Tensor,
+        swa_cache: torch.Tensor,
+        swa_indices: torch.Tensor,
+        swa_lens: torch.Tensor,
+        extra_cache: torch.Tensor | None,
+        extra_indices: torch.Tensor | None,
+        extra_lens: torch.Tensor | None,
+        output: torch.Tensor,
+    ) -> None:
+        """One DCP attention step for ``q``/``output`` of ``[T, n_local_heads, 512]``.
+
+        Gather the local heads to ``[T, n_local_heads * W, 512]`` (TP order),
+        run this rank's sparse attention over its compressed-state shard with
+        the LSE, then combine across the DCP group and reduce-scatter the heads
+        back. The replicated SWA window and the sink are attended on DCP rank 0
+        only, so with ``L_r`` the per-rank base-2 LSEs (sink folded into
+        ``L_0``) the combine ``L = log2(sum_r 2^L_r)`` and ``o = sum_r
+        2^(L_r - L) o_r`` reproduce the DCP1 softmax exactly.
+        """
+        assert self.dcp_group is not None
+        num_tokens = q.shape[0]
+        # Timing events cannot be synchronised inside a CUDA-graph capture (the
+        # eager break runs on the capture stream during warmup/capture).
+        capturing = torch.cuda.is_current_stream_capturing()
+        prefill_timing = _DSV41_PREFILL_TIMING and num_tokens >= 1024 and not capturing
+        timing = (self._dcp_timing or prefill_timing) and not capturing
+        if self._dcp_timing:
+            logger.info_once(
+                "DCP path active on layer %d (capturing=%s, T=%d)",
+                self.layer_id,
+                capturing,
+                num_tokens,
+            )
+        if timing:
+            ev = [torch.cuda.Event(enable_timing=True) for _ in range(4)]
+            ev[0].record()
+        # Collective: every rank builds the gathered sink (lazily, once); only
+        # DCP rank 0 applies it, together with the replicated SWA window.
+        sink = self._dcp_attn_sink()
+        q_gathered = self.dcp_group.all_gather(q, dim=1)
+        num_heads = q_gathered.shape[1]
+        if timing:
+            ev[1].record()
+        out_gathered = torch.empty(
+            (num_tokens, num_heads, self.head_dim), dtype=torch.bfloat16, device=q.device
+        )
+        lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=q.device)
+        if self.dcp_rank != 0:
+            swa_indices, swa_lens = self._dcp_neutral_swa(
+                num_tokens, swa_indices.shape[-1], q.device
+            )
+        _sm120_sparse_attention_with_lse(
+            q_gathered,
+            swa_cache,
+            self._get_workspace(q.device),
+            swa_indices,
+            swa_lens,
+            extra_cache,
+            extra_indices,
+            extra_lens,
+            out_gathered,
+            self.scale,
+            sink if self.dcp_rank == 0 else None,
+            lse,
+        )
+        if timing:
+            ev[2].record()
+        # Rows with no local candidate (and no sink) carry the kernel's -1e30
+        # LSE and a zero output; the combine kernel's max-subtracted exp2 makes
+        # that the merge identity without any masking (test_stage3_combine).
+        if self._dcp_combine_a2a:
+            combined = dcp_a2a_lse_reduce(
+                out_gathered, lse, self.dcp_group, is_lse_base_on_e=False
+            )
+        else:
+            combined = cp_lse_ag_out_rs(
+                out_gathered,
+                lse,
+                self.dcp_group,
+                ctx=self._cp_ctx,
+                is_lse_base_on_e=False,
+            )
+        output.copy_(combined)
+        if timing:
+            ev[3].record()
+            ev[3].synchronize()
+            if prefill_timing:
+                _dsv41_prefill_acc_add("attn_gather", ev[0].elapsed_time(ev[1]))
+                _dsv41_prefill_acc_add("attn_kernel", ev[1].elapsed_time(ev[2]))
+                _dsv41_prefill_acc_add("attn_combine", ev[2].elapsed_time(ev[3]))
+            acc = self._dcp_t_acc
+            acc[0] += ev[0].elapsed_time(ev[1])
+            acc[1] += ev[1].elapsed_time(ev[2])
+            acc[2] += ev[2].elapsed_time(ev[3])
+            acc[3] += 1
+            if self._dcp_timing and acc[3] % 100 == 0 and self.dcp_rank == 0:
+                logger.info(
+                    "DCP timing layer %d (T=%d): gather %.3f ms, kernel %.3f ms, "
+                    "combine %.3f ms per call (avg of %d)",
+                    self.layer_id,
+                    num_tokens,
+                    acc[0] / acc[3],
+                    acc[1] / acc[3],
+                    acc[2] / acc[3],
+                    acc[3],
+                )
+                acc[0] = acc[1] = acc[2] = 0.0
+                acc[3] = 0
+
     def _forward_decode(
         self,
         q: torch.Tensor,
@@ -789,10 +1196,11 @@
                 raise RuntimeError(
                     "Compressed-layer decode requires top-k indices from the indexer."
                 )
-            # Local indices filled by the index-source layer's indexer.
+            # Indices filled by the index-source layer's indexer (request-local
+            # at DCP1, global compressed ids under DCP).
             is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
             block_size = attn_metadata.block_size // self.compress_ratio
-            global_indices, extra_sparse_lengths = compute_global_topk_indices_and_lens(
+            global_indices, extra_sparse_lengths = self._compressed_topk_to_slots(
                 self.topk_indices_buffer[:num_decode_tokens],
                 swa_metadata.token_to_req_indices,
                 attn_metadata.block_table[:num_decodes],
@@ -812,6 +1220,20 @@
             raise RuntimeError(
                 "Compressed sparse MLA decode requires compressed sparse indices."
             )
+        if self.dcp_world_size > 1 and not swa_only:
+            # SWA-only layers keep the DCP1 path: their cache is replicated and
+            # every rank runs its own heads over the whole window.
+            self._run_sm120_dcp(
+                q,
+                swa_cache,
+                swa_indices,
+                swa_lens,
+                extra_cache,
+                extra_sparse_indices,
+                extra_sparse_lengths,
+                output,
+            )
+            return
         flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
             query=q,
             swa_kv_cache=swa_cache,
@@ -849,6 +1271,7 @@
 
         extra_sparse_indices: torch.Tensor | None = None
         extra_sparse_lengths: torch.Tensor | None = None
+        gathered_kv: torch.Tensor | None = None
         if not swa_only:
             if self.topk_indices_buffer is None:
                 raise RuntimeError(
@@ -872,15 +1295,30 @@
                 num_decode_tokens, num_decode_tokens + num_prefill_tokens
             )
             block_size = attn_metadata.block_size // self.compress_ratio
-            extra_sparse_indices, extra_sparse_lengths = (
-                compute_global_topk_indices_and_lens(
-                    local_topk_indices,
-                    swa_metadata.token_to_req_indices[prefill_token_slice],
-                    attn_metadata.block_table,
-                    block_size,
-                    swa_metadata.is_valid_token[prefill_token_slice],
+            if self.dcp_world_size > 1 and _DCP_PREFILL_GATHER:
+                assert compressed_k_cache is not None
+                gathered_kv, extra_sparse_indices, extra_sparse_lengths = (
+                    self._dcp_prefill_gather(
+                        compressed_k_cache,
+                        local_topk_indices,
+                        swa_metadata.token_to_req_indices[prefill_token_slice],
+                        swa_metadata.is_valid_token[prefill_token_slice],
+                        attn_metadata,
+                        num_decodes,
+                        num_prefills,
+                        block_size,
+                    )
+                )
+            else:
+                extra_sparse_indices, extra_sparse_lengths = (
+                    self._compressed_topk_to_slots(
+                        local_topk_indices,
+                        swa_metadata.token_to_req_indices[prefill_token_slice],
+                        attn_metadata.block_table,
+                        block_size,
+                        swa_metadata.is_valid_token[prefill_token_slice],
+                    )
                 )
-            )
 
         assert swa_metadata.prefill_swa_indices is not None
         assert swa_metadata.prefill_swa_lens is not None
@@ -894,7 +1332,11 @@
                 raise RuntimeError(
                     "Compressed sparse MLA layers require their compressed KV cache."
                 )
-            extra_kv_paged = self._as_sparse_cache(compressed_k_cache)
+            extra_kv_paged = (
+                gathered_kv
+                if gathered_kv is not None
+                else self._as_sparse_cache(compressed_k_cache)
+            )
 
         num_chunks = (
             num_prefills + self.PREFILL_CHUNK_SIZE - 1
@@ -927,17 +1369,40 @@
                 raise RuntimeError(
                     "Compressed sparse MLA prefill requires compressed sparse indices."
                 )
-            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
-                query=q_chunk,
-                swa_kv_cache=swa_kv_paged,
-                workspace_buffer=self._get_workspace(q.device),
-                sparse_indices=swa_indices_chunk,
-                compressed_kv_cache=extra_kv_paged,
-                out=output[query_start:query_end],
-                bmm1_scale=self.scale,
-                sinks=self.attn_sink,
-                kv_layout="NHD",
-                swa_topk_lens=swa_lens_chunk,
-                extra_sparse_indices=extra_sparse_indices_chunk,
-                extra_sparse_topk_lens=extra_sparse_lengths_chunk,
-            )
+            if self.dcp_world_size > 1 and not swa_only and gathered_kv is None:
+                # Chunk bounds come from the CPU query_start_loc (identical on
+                # every rank), so the collectives below line up across ranks.
+                qs, qe = int(query_start), int(query_end)
+                for sub_start in range(qs, qe, self.DCP_PREFILL_TOKEN_CHUNK):
+                    sub_end = min(sub_start + self.DCP_PREFILL_TOKEN_CHUNK, qe)
+                    assert extra_sparse_indices is not None
+                    assert extra_sparse_lengths is not None
+                    self._run_sm120_dcp(
+                        q[sub_start:sub_end],
+                        swa_kv_paged,
+                        swa_metadata.prefill_swa_indices[sub_start:sub_end],
+                        swa_metadata.prefill_swa_lens[sub_start:sub_end],
+                        extra_kv_paged,
+                        extra_sparse_indices[sub_start:sub_end],
+                        extra_sparse_lengths[sub_start:sub_end],
+                        output[sub_start:sub_end],
+                    )
+                continue
+            with _DSV41PrefillTimer(
+                "attn_kernel",
+                int(query_end - query_start) if gathered_kv is not None else 0,
+            ):
+                flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
+                    query=q_chunk,
+                    swa_kv_cache=swa_kv_paged,
+                    workspace_buffer=self._get_workspace(q.device),
+                    sparse_indices=swa_indices_chunk,
+                    compressed_kv_cache=extra_kv_paged,
+                    out=output[query_start:query_end],
+                    bmm1_scale=self.scale,
+                    sinks=self.attn_sink,
+                    kv_layout="NHD",
+                    swa_topk_lens=swa_lens_chunk,
+                    extra_sparse_indices=extra_sparse_indices_chunk,
+                    extra_sparse_topk_lens=extra_sparse_lengths_chunk,
+                )