launcher/dcp2/diffs/04-model_executor_layers_sparse_attn_indexer.py.diff

--- baked-ours/model_executor/layers/sparse_attn_indexer.py	2026-09-10 20:13:20.000000000 +0200
+++ baked/model_executor/layers/sparse_attn_indexer.py	2026-09-11 10:13:47.000000000 +0200
@@ -17,6 +17,9 @@
     apply_candidate_mask as _apply_candidate_mask,
 )
 from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+    _PrefillTimer as _DSV41PrefillTimer,
+)
+from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
     select_candidate_blocks as _select_candidate_blocks,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -325,6 +328,7 @@
     candidate_blocks: torch.Tensor | None = None,
     candidate_block_size: int = 0,
     candidate_write: bool = False,
+    compress_ratio: int = 1,
 ) -> torch.Tensor:
     # careful! this will be None in dummy run
     forward_context = get_forward_context()
@@ -332,13 +336,25 @@
     fp8_dtype = current_platform.fp8_dtype()
     k_cache_prefix = _resolve_layer_name(k_cache_prefix)
 
+    # v4.1 two-level candidate filtering under DCP: the logits columns are this
+    # rank's local compressed states; with interleave 1 and world | block, local
+    # block b (size block // world) is exactly this rank's shard of global block
+    # b, so candidate ids need no translation. The producer MAX-reduces block
+    # scores across ranks and pins the global newest block (candidate_blocks.py).
+    local_cand_bs = candidate_block_size
+    dcp_group = None
     if candidate_blocks is not None:
-        # Candidate blocks are request-local; the DCP-sharded logits layout
-        # would need per-rank translation that is not implemented.
-        assert dcp_world_size == 1, (
-            "v4.1 two-level candidate filtering is not supported with DCP."
-        )
         assert candidate_block_size > 0
+        if dcp_world_size > 1:
+            assert cp_kv_cache_interleave_size == 1, (
+                "v4.1 candidate filtering under DCP needs cp_kv_cache_interleave_size=1"
+            )
+            assert candidate_block_size % dcp_world_size == 0, (
+                f"candidate_block_size ({candidate_block_size}) must be a multiple "
+                f"of dcp_world_size ({dcp_world_size})"
+            )
+            local_cand_bs = candidate_block_size // dcp_world_size
+            dcp_group = get_dcp_group()
 
     # assert isinstance(attn_metadata, dict)
     if not isinstance(attn_metadata, dict):
@@ -494,6 +510,23 @@
             if chunk.local_total_seq_lens == 0:
                 logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
                 topk_indices.fill_(-1)
+                if candidate_blocks is not None and candidate_write and dcp_group is not None:
+                    # Other ranks may hold states for this chunk: join the MAX
+                    # reduce with an empty local contribution.
+                    _select_candidate_blocks(
+                        logits,
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                        candidate_blocks.shape[1],
+                        local_cand_bs,
+                        candidate_blocks[chunk.token_start : chunk.token_end],
+                        dcp_group=dcp_group,
+                        global_row_ke=chunk.global_ctx_lens,
+                        global_block_size=candidate_block_size,
+                        max_global_blocks=triton.cdiv(
+                            chunk.max_global_ctx, candidate_block_size
+                        ),
+                    )
             else:
                 # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
                 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
@@ -539,8 +572,16 @@
                             cu_seqlen_ks,
                             cu_seqlen_ke,
                             chunk_candidates.shape[1],
-                            candidate_block_size,
+                            local_cand_bs,
                             chunk_candidates,
+                            dcp_group=dcp_group,
+                            global_row_ke=chunk.global_ctx_lens,
+                            global_block_size=candidate_block_size,
+                            max_global_blocks=(
+                                triton.cdiv(chunk.max_global_ctx, candidate_block_size)
+                                if dcp_group is not None
+                                else None
+                            ),
                         )
                     else:
                         _apply_candidate_mask(
@@ -548,7 +589,7 @@
                             cu_seqlen_ks,
                             cu_seqlen_ke,
                             chunk_candidates,
-                            candidate_block_size,
+                            local_cand_bs,
                         )
                 ops.top_k_per_row_prefill(
                     logits,
@@ -666,14 +707,44 @@
             vis = vis[:num_rows]
             decode_candidates = candidate_blocks[:num_rows]
             if candidate_write:
+                global_row_ke = None
+                max_global_blocks = None
+                if dcp_group is not None:
+                    # Global (un-sharded) per-row context in STATES for the
+                    # newest-block pin: per token when flattened / next_n == 1,
+                    # per request (row_repeat) or expanded per row otherwise.
+                    gl = decode_metadata.global_seq_lens
+                    assert gl is not None
+                    gl = gl.reshape(-1)
+                    if row_repeat > 1:
+                        gl = gl[:batch_size]
+                    elif gl.numel() != num_rows:
+                        # native spec decode: request lens -> per-row lens
+                        gl = (
+                            gl[:batch_size].unsqueeze(1)
+                            - next_n
+                            + 1
+                            + torch.arange(next_n, device=gl.device, dtype=gl.dtype)
+                        ).clamp_(min=0).reshape(-1)
+                    else:
+                        gl = gl[:num_rows]
+                    global_row_ke = gl // compress_ratio
+                    max_global_blocks = triton.cdiv(
+                        attn_metadata_narrowed.max_seq_len // compress_ratio,
+                        candidate_block_size,
+                    )
                 _select_candidate_blocks(
                     logits,
                     None,
                     vis,
                     decode_candidates.shape[1],
-                    candidate_block_size,
+                    local_cand_bs,
                     decode_candidates,
                     row_repeat,
+                    dcp_group=dcp_group,
+                    global_row_ke=global_row_ke,
+                    global_block_size=candidate_block_size,
+                    max_global_blocks=max_global_blocks,
                 )
             else:
                 _apply_candidate_mask(
@@ -681,7 +752,7 @@
                     None,
                     vis,
                     decode_candidates,
-                    candidate_block_size,
+                    local_cand_bs,
                     row_repeat,
                 )
         topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
@@ -790,6 +861,7 @@
     candidate_blocks: torch.Tensor | None = None,
     candidate_block_size: int = 0,
     candidate_write: bool = False,
+    compress_ratio: int = 1,
 ) -> torch.Tensor:
     return topk_indices_buffer
 
@@ -904,6 +976,17 @@
         k: torch.Tensor | None,
         weights: torch.Tensor,
     ):
+        # Stage 7: phase "indexer" of the prefill timing (inert unless enabled).
+        with _DSV41PrefillTimer("indexer", int(hidden_states.shape[0])):
+            return self._forward_cuda_impl(hidden_states, q_quant, k, weights)
+
+    def _forward_cuda_impl(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor | None,
+        weights: torch.Tensor,
+    ):
         # FP8 path: single tensor (per-token scale is folded into `weights`).
         # FP4 path: (values, scales) tuple with scales required by the kernel.
         if isinstance(q_quant, tuple):
@@ -935,6 +1018,7 @@
             candidate_blocks=self.candidate_blocks,
             candidate_block_size=self.candidate_block_size,
             candidate_write=self.candidate_write,
+            compress_ratio=self.compress_ratio,
         )
 
     def forward_xpu(