launcher/dcp2/diffs/02-model_executor_kernels_attention_dsa_candidate_blocks.py.diff

--- baked-ours/model_executor/kernels/attention/dsa/candidate_blocks.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/model_executor/kernels/attention/dsa/candidate_blocks.py	2026-09-11 10:13:47.000000000 +0200
@@ -1,11 +1,67 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 import torch
 
 from vllm.triton_utils import tl, triton
 
 
+def _cpu_kernels_allowed() -> bool:
+    # CPU tests run the Triton kernels through the interpreter.
+    return os.environ.get("TRITON_INTERPRET", "0") == "1"
+
+
+# Scratch cap for the DCP candidate-score all-reduce (bytes per slice).
+_DCP_REDUCE_BYTES = 128 * 1024 * 1024
+
+# Stage 7: opt-in prefill phase timing (DSV41_PREFILL_TIMING=1). Phases add GPU
+# milliseconds here; the runner reports and clears them after each eager step
+# of >= 1024 tokens. Inert (no events, no syncs) unless enabled.
+_PREFILL_TIMING = os.environ.get("DSV41_PREFILL_TIMING", "0") == "1"
+PREFILL_ACC: dict[str, float] = {}
+
+
+def prefill_acc_add(key: str, ms: float) -> None:
+    PREFILL_ACC[key] = PREFILL_ACC.get(key, 0.0) + ms
+
+
+def prefill_acc_take() -> dict[str, float]:
+    d = dict(PREFILL_ACC)
+    PREFILL_ACC.clear()
+    return d
+
+
+class _PrefillTimer:
+    """`with _PrefillTimer(key, num_tokens): ...` — CUDA events around the block,
+    synchronised at exit (eager prefill steps only; never while capturing)."""
+
+    __slots__ = ("key", "on", "e0", "e1")
+
+    def __init__(self, key: str, num_tokens: int) -> None:
+        self.key = key
+        self.on = (
+            _PREFILL_TIMING
+            and num_tokens >= 1024
+            and not torch.cuda.is_current_stream_capturing()
+        )
+
+    def __enter__(self):
+        if self.on:
+            self.e0 = torch.cuda.Event(enable_timing=True)
+            self.e1 = torch.cuda.Event(enable_timing=True)
+            self.e0.record()
+        return self
+
+    def __exit__(self, *exc):
+        if self.on:
+            self.e1.record()
+            self.e1.synchronize()
+            prefill_acc_add(self.key, self.e0.elapsed_time(self.e1))
+        return False
+
+
 @triton.jit
 def _max_with_nan(a, b):
     return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL)
@@ -27,6 +83,7 @@
     HAS_STARTS: tl.constexpr,
     ROW_REPEAT: tl.constexpr,
     TILE: tl.constexpr,
+    PIN_NEWEST: tl.constexpr = True,
 ):
     row = tl.program_id(0).to(tl.int64)
     blocks = tl.program_id(1) * TILE + tl.arange(0, TILE)
@@ -43,11 +100,12 @@
         other=-float("inf"),
     )
     reduced = tl.reduce(values, 1, _max_with_nan)
-    reduced = tl.where(
-        (end > start) & (blocks == (end - start - 1) // BLOCK_SIZE),
-        float("inf"),
-        reduced,
-    )
+    if PIN_NEWEST:
+        reduced = tl.where(
+            (end > start) & (blocks == (end - start - 1) // BLOCK_SIZE),
+            float("inf"),
+            reduced,
+        )
     tl.store(scores + row * nblocks + blocks, reduced, blocks < nblocks)
 
 
@@ -137,27 +195,15 @@
     )
 
 
-def select_candidate_blocks(
+def _block_scores(
     logits: torch.Tensor,
     row_ks: torch.Tensor | None,
     row_ke: torch.Tensor,
-    topk_blocks: int,
     block_size: int,
-    out: torch.Tensor,
-    row_repeat: int = 1,
-) -> None:
-    """Select local block IDs by maximum score, pinning each row's newest block.
-
-    Row bounds are in packed column space; absent starts mean zero.
-    Decode rows share bounds in groups of ``row_repeat``. Output is -1 padded.
-    """
-    assert logits.is_cuda
+    row_repeat: int,
+    pin_newest: bool,
+) -> torch.Tensor:
     rows, width = logits.shape
-    if not rows:
-        return
-    if not width:
-        out.fill_(-1)
-        return
     nblocks = triton.cdiv(width, block_size)
     scores = logits.new_empty((rows, nblocks))
     _block_scores_kernel[(rows, triton.cdiv(nblocks, 128))](
@@ -174,7 +220,109 @@
         row_ks is not None,
         row_repeat,
         128,
+        PIN_NEWEST=pin_newest,
     )
+    return scores
+
+
+def select_candidate_blocks(logits: torch.Tensor, *args, **kwargs) -> None:
+    """Stage 7: timed wrapper (phase "candidates"); see _select_candidate_blocks_impl."""
+    with _PrefillTimer("candidates", int(logits.shape[0])):
+        return _select_candidate_blocks_impl(logits, *args, **kwargs)
+
+
+def _select_candidate_blocks_impl(
+    logits: torch.Tensor,
+    row_ks: torch.Tensor | None,
+    row_ke: torch.Tensor,
+    topk_blocks: int,
+    block_size: int,
+    out: torch.Tensor,
+    row_repeat: int = 1,
+    *,
+    dcp_group=None,
+    global_row_ke: torch.Tensor | None = None,
+    global_block_size: int | None = None,
+    max_global_blocks: int | None = None,
+) -> None:
+    """Select local block IDs by maximum score, pinning each row's newest block.
+
+    Row bounds are in packed column space; absent starts mean zero.
+    Decode rows share bounds in groups of ``row_repeat``. Output is -1 padded.
+
+    Under DCP (``dcp_group`` given) ``logits`` hold this rank's local packed
+    columns and ``block_size`` is the LOCAL block, ``global_block_size //
+    dcp_world_size``: with interleave 1, local block ``b`` is exactly this
+    rank's states of global block ``b``, so the global block score is the MAX
+    over ranks of the local block scores (one all-reduce over
+    ``[rows, max_global_blocks]``) and the block ids need no translation. The
+    newest block is pinned from the GLOBAL per-row context ``global_row_ke``
+    (states, per row or per ``row_repeat`` group) after the reduce, so every
+    rank pins the same block. Collective: every DCP rank must call this with
+    the same ``rows`` and ``max_global_blocks`` (an empty local shard is fine).
+    """
+    assert logits.is_cuda or _cpu_kernels_allowed()
+    rows, width = logits.shape
+    if not rows:
+        return
+    if dcp_group is None:
+        if not width:
+            out.fill_(-1)
+            return
+        scores = _block_scores(logits, row_ks, row_ke, block_size, row_repeat, True)
+        nblocks = scores.shape[1]
+    else:
+        assert global_row_ke is not None
+        assert global_block_size is not None and max_global_blocks is not None
+        assert global_block_size == block_size * dcp_group.world_size, (
+            "DCP candidate blocks need block_size == global_block_size // world"
+        )
+        nblocks = max(int(max_global_blocks), 1)
+        ends = global_row_ke.reshape(-1)
+        if row_repeat > 1:
+            ends = ends.repeat_interleave(row_repeat)
+        ends = ends[:rows].to(torch.int64)
+        # Pin each row's newest GLOBAL block (DCP1 pins (end - start - 1) //
+        # block from the global bounds); rows without context select nothing.
+        pin = ((ends - 1) // global_block_size).clamp_(0, nblocks - 1).view(-1, 1)
+        local = (
+            _block_scores(logits, row_ks, row_ke, block_size, row_repeat, False)
+            if width
+            else None
+        )
+        # Reduce in row slices: a 16K-token prefill chunk at 128K context would
+        # otherwise need a ~1 GiB scratch per rank (GB10 unified memory is the
+        # binding constraint). Slice count derives from rows/nblocks, which are
+        # identical on every rank, so the collectives still line up.
+        k = min(topk_blocks, nblocks)
+        rows_per_slice = max(1, _DCP_REDUCE_BYTES // (nblocks * 4))
+        top_values = logits.new_empty((rows, k))
+        top_indices = torch.empty((rows, k), dtype=torch.int64, device=logits.device)
+        for r0 in range(0, rows, rows_per_slice):
+            r1 = min(r0 + rows_per_slice, rows)
+            scores = logits.new_full((r1 - r0, nblocks), -float("inf"))
+            if local is not None:
+                n = min(local.shape[1], nblocks)
+                scores[:, :n] = local[r0:r1, :n]
+            torch.distributed.all_reduce(
+                scores, op=torch.distributed.ReduceOp.MAX, group=dcp_group.device_group
+            )
+            scores.scatter_(1, pin[r0:r1], float("inf"))
+            scores.masked_fill_((ends[r0:r1] <= 0).view(-1, 1), -float("inf"))
+            # Keep the existing top-k tie behavior.
+            sl = scores.topk(k, dim=-1)
+            top_values[r0:r1] = sl.values
+            top_indices[r0:r1] = sl.indices
+        _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](
+            top_values,
+            top_indices,
+            out,
+            *out.stride(),
+            k,
+            topk_blocks,
+            256,
+        )
+        return
     # Keep the existing top-k tie behavior.
     top = scores.topk(min(topk_blocks, nblocks), dim=-1)
     _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](
@@ -196,8 +344,13 @@
     block_size: int,
     row_repeat: int = 1,
 ) -> None:
-    """Mask packed logits outside causal bounds and request-local candidates."""
-    assert logits.is_cuda
+    """Mask packed logits outside causal bounds and request-local candidates.
+
+    Under DCP pass the LOCAL block size (global // world): local block ``b``
+    is this rank's shard of global block ``b``, so the candidate ids apply
+    unchanged.
+    """
+    assert logits.is_cuda or _cpu_kernels_allowed()
     rows, width = logits.shape
     if not rows or not width:
         return