launcher/dcp2/diffs/14-v1_attention_backends_mla_indexer.py.diff

--- baked-ours/v1/attention/backends/mla/indexer.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/v1/attention/backends/mla/indexer.py	2026-09-11 10:13:47.000000000 +0200
@@ -314,6 +314,11 @@
     local_cu_seq_lens: torch.Tensor | None = None
     local_total_seq_lens: int = 0
     max_local_total_seq_lens: int = 0
+    # Per query token: GLOBAL compressed context length (states), i.e. the
+    # un-sharded `cu_seqlen_ke - cu_seqlen_ks`. Under DCP the candidate-block
+    # selection pins each row's newest global block from this.
+    global_ctx_lens: torch.Tensor | None = None
+    max_global_ctx: int = 0
 
 
 class BuildPrefillChunkMetadataKernel(
@@ -346,6 +351,8 @@
         token_to_seq_ptr,
         cu_compressed_seq_len_ks_ptr,
         cu_compressed_seq_len_ke_ptr,
+        # [output_query_len] global (un-sharded) compressed context per token
+        global_ctx_len_ptr,
         query_slice_start,
         query_slice_stop,
         DCP_RANK,
@@ -388,6 +395,7 @@
             # global per-token length is sharded across ranks.
             global_ctx = start_pos + 1 + offset
             len_per_token = global_ctx // COMPRESS_RATIO
+            tl.store(global_ctx_len_ptr + out_pos, len_per_token, mask=mask)
             if DCP_WORLD > 1:
                 # Per-rank local context length under interleave-aware DCP, matching
                 # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
@@ -478,6 +486,7 @@
             token_to_seq=int32_ptr,
             cu_compressed_seq_len_ks=int32_ptr,
             cu_compressed_seq_len_ke=int32_ptr,
+            global_ctx_len=int32_ptr,
             query_slice_start=compile_key.query_slice_start,
             query_slice_stop=compile_key.query_slice_stop,
             DCP_RANK=compile_key.dcp_rank,
@@ -497,6 +506,7 @@
         token_to_seq: torch.Tensor,
         cu_compressed_seq_len_ks: torch.Tensor,
         cu_compressed_seq_len_ke: torch.Tensor,
+        global_ctx_len: torch.Tensor,
         query_slice_start: int,
         query_slice_stop: int,
         DCP_RANK: int,
@@ -697,6 +707,9 @@
     # so this is None; its own split uses self.decode_threshold.
     reorder_batch_threshold: int | None = None
     requires_block_table_width = True
+    # DeepSeek-V4.1 DSpark under DCP: the draft is SWA-only (no indexer); the
+    # target's non-causal verification rows use the DCP indexer path.
+    supports_non_causal_multi_token_dcp = True
 
     @classmethod
     def get_cudagraph_support(
@@ -812,9 +825,14 @@
             assert isinstance(self.kv_cache_spec.tokens_per_state, int)
             self.compress_ratio = self.kv_cache_spec.tokens_per_state
         if self.dcp_world_size > 1 and self.compress_ratio > 1:
-            raise NotImplementedError(
-                "DCP is not supported with sparse indexer KV compression "
-                f"(compress_ratio={self.compress_ratio})."
+            # DCP + compressed indexer cache (DeepSeek-V4.1 ratio-2 index
+            # sources): the K cache is sharded over compressed STATES; the
+            # compressed slot mapping and the decode bounds below are computed
+            # in state coordinates (DCP2-DESIGN-stage23.md §A.1).
+            logger.info_once(
+                "DSA indexer: DCP%d with compress_ratio=%d (state-sharded K cache)",
+                self.dcp_world_size,
+                self.compress_ratio,
             )
 
         # Pre-allocate buffers for CUDA graph compatibility when
@@ -1057,6 +1075,10 @@
                 and self.kv_cache_spec.block_size % kernel_block_size == 0
             ):
                 factor = self.kv_cache_spec.block_size // kernel_block_size
+                assert self.dcp_world_size == 1, (
+                    "DCP with a compressed indexer cache needs kernel_block_size == "
+                    f"kv_cache_spec.block_size (got factor {factor})."
+                )
                 indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
             padded_num_tokens = num_tokens
             if self.pcp_world_size > 1:
@@ -1069,6 +1091,9 @@
                 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,
             )
             if self.pcp_world_size > 1:
                 compressed_slot_mapping = get_pcp_group().all_gather(
@@ -1235,16 +1260,12 @@
             # batches. Keep its address stable across varlen graph replays.
             seq_lens_is_buffer_view = not use_native or next_n > 1
 
-            # DCP: localize the now-expanded per-token global bounds to this
-            # rank's owned KV. Done here (after expansion) so each token's global
-            # causal length is localized individually; see the comment above.
-            if dcp_local_seq_lens is not None:
-                seq_lens = self._dcp_localize_decode_seq_lens(
-                    seq_lens, num_decodes, seq_lens_is_buffer_view
-                )
-
             # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
             # compressed tokens. Convert uncompressed seq_lens to compressed.
+            # This must happen BEFORE the DCP localization below: the cache is
+            # sharded over compressed states, so this rank's share is
+            # localize(len // R), not localize(len) // R (len=10, R=2, W=2:
+            # rank 0 owns states {0, 2, 4} -> 3, but 5 // 2 == 2).
             if self.compress_ratio > 1:
                 if seq_lens_is_buffer_view:
                     seq_lens //= self.compress_ratio
@@ -1255,6 +1276,17 @@
                     )
                     self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
                     seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
+                    # The copy is private: localize it in place below.
+                    seq_lens_is_buffer_view = True
+
+            # DCP: localize the now-expanded (and now-compressed) per-token
+            # global bounds to this rank's owned KV. Done here (after expansion)
+            # so each token's global causal length is localized individually;
+            # see the comment above.
+            if dcp_local_seq_lens is not None:
+                seq_lens = self._dcp_localize_decode_seq_lens(
+                    seq_lens, num_decodes, seq_lens_is_buffer_view
+                )
 
             # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
             # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
@@ -1367,6 +1399,7 @@
 
     cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
     cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
+    global_ctx_len = torch.empty(output_query_len, dtype=torch.int32, device=device)
 
     # Under DCP the kernel writes this rank's local row bounds into
     # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
@@ -1378,6 +1411,7 @@
         token_to_seq,
         cu_seq_len_ks,
         cu_seq_len_ke,
+        global_ctx_len,
         qs_start,
         qs_stop,
         dcp_rank,
@@ -1409,6 +1443,12 @@
         local_cu_seq_lens=local_cu_seq_lens,
         local_total_seq_lens=local_total_seq_lens,
         max_local_total_seq_lens=max_local_total_seq_lens,
+        global_ctx_lens=global_ctx_len if dcp_world_size > 1 else None,
+        max_global_ctx=(
+            int(compressed_seq_lens_cpu[start_idx:end_idx].max().item())
+            if dcp_world_size > 1
+            else 0
+        ),
     )