--- 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