launcher/dcp2/diffs/11-models_deepseek_v4_1_sparse_mla.py.diff

--- baked-ours/models/deepseek_v4_1/sparse_mla.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/models/deepseek_v4_1/sparse_mla.py	2026-09-11 10:13:47.000000000 +0200
@@ -20,14 +20,18 @@
     MultipleOf,
 )
 from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.attention.backends.mla.sparse_swa import (
     _LAYER_TYPE_C1A,
     _LAYER_TYPE_C2A,
     _LAYER_TYPE_SWAONLY,
     DeepseekSparseSWAMetadataBuilder,
 )
+from vllm.logger import init_logger
 from vllm.v1.kv_cache_interface import AttentionSpec
 
+logger = init_logger(__name__)
+
 # v4.1 per-layer compress ratios: 0 = pure sliding window, 1 = full-length
 # compressed cache, 2 = ratio-2 compressed cache. Ratio-1 and ratio-2 layers
 # both attend over indexer topk indices into a shared compressed cache but
@@ -58,6 +62,8 @@
     tile-scheduler plans from the v4.1 topology.
     """
 
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = True
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         compress_ratios = getattr(
@@ -139,12 +145,23 @@
     req_id_per_token: torch.Tensor
     block_size: int
     topk_tokens: int
+    # DCP: the compressed cache is sharded over STATES with this interleave;
+    # block_size stays the per-rank kernel page (tokens). seq_lens are the
+    # global lengths (tokens); dcp_local_seq_lens this rank's share of the
+    # compressed length (states), None at DCP1.
+    cp_kv_cache_interleave_size: int = 1
+    seq_lens: torch.Tensor | None = None
+    dcp_local_seq_lens: torch.Tensor | None = None
 
 
 class DeepseekV4SparseMLAMetadataBuilder(
     AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata]
 ):
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    # The SM120 backend gathers queries / combines per token, so DSpark's
+    # non-causal multi-token decode is supported under DCP (the draft layers
+    # themselves are SWA-only and bypass the DCP path).
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = True
 
     def __init__(
         self,
@@ -155,9 +172,22 @@
     ) -> None:
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
         self.model_config = vllm_config.model_config
+        # DCP geometry (self.dcp_world_size / dcp_rank come from the base
+        # builder: process DCP for this sharded group). The V4.1 SM120 backend
+        # gathers queries and combines LSEs per token, so multi-token (DSpark)
+        # decodes stay decodes under DCP.
+        self.cp_kv_cache_interleave_size = (
+            vllm_config.parallel_config.cp_kv_cache_interleave_size
+        )
+        if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size != 1:
+            raise NotImplementedError(
+                "DeepSeek-V4.1 DCP supports cp_kv_cache_interleave_size=1 only"
+            )
         # Classify single-token queries (plus num_speculative_tokens via
         # supports_spec_as_decode=True) as decodes; longer queries go to prefill.
-        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
+        self._init_reorder_batch_threshold(
+            1, supports_spec_as_decode=True, supports_dcp_with_varlen=True
+        )
         self.topk_tokens = self.model_config.hf_config.index_topk
 
         max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -174,6 +204,20 @@
                 "DeepSeek V4.1 compressed-KV caches use compress_ratio 1 or 2; "
                 f"got {self.compress_ratio}."
             )
+        if self.dcp_world_size > 1:
+            # Per-rank page geometry the state-space sharding relies on: the
+            # builder must see the kernel page (block_size tokens = num_states
+            # states per rank), never the manager block (x dcp).
+            assert self.kv_cache_spec.block_size % self.compress_ratio == 0
+            logger.info_once(
+                "DeepSeek-V4.1 compressed KV under DCP%d: ratio %d, %d tokens "
+                "(%d states) per page per rank, interleave %d",
+                self.dcp_world_size,
+                self.compress_ratio,
+                self.kv_cache_spec.block_size,
+                int(self.kv_cache_spec.num_states),
+                self.cp_kv_cache_interleave_size,
+            )
 
         # Pre-allocate compressed slot mapping buffer for CUDA graph address
         # stability when compress_ratio > 1.
@@ -201,6 +245,17 @@
                 int(self.kv_cache_spec.num_states),
                 self.compress_ratio,
                 out=self.compressed_slot_mapping_buffer,
+                dcp_world_size=self.dcp_world_size,
+                dcp_rank=self.dcp_rank,
+                cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            )
+        dcp_local_seq_lens = None
+        if self.dcp_world_size > 1:
+            dcp_local_seq_lens = get_dcp_local_seq_lens(
+                cm.seq_lens // self.compress_ratio,
+                self.dcp_world_size,
+                self.dcp_rank,
+                self.cp_kv_cache_interleave_size,
             )
 
         return DeepseekV4FlashMLAMetadata(
@@ -214,6 +269,9 @@
             req_id_per_token=req_id_per_token,
             block_size=self.kv_cache_spec.block_size,
             topk_tokens=self.topk_tokens,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            seq_lens=cm.seq_lens,
+            dcp_local_seq_lens=dcp_local_seq_lens,
         )