diff --git a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py index b9cdc923b1..287be45df1 100644 --- a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py +++ b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py @@ -3,8 +3,13 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +_IDX_BITS = tl.constexpr(20) +_IDX_MASK = tl.constexpr((1 << 20) - 1) +_NEG_INF = tl.constexpr(-float("inf")) + @triton.jit def _max_with_nan(a, b): @@ -138,6 +143,252 @@ def _mask_candidates_kernel( ) +@triton.jit +def _pack_key(score, index, valid): + """Order-preserving int64 key: float order << _IDX_BITS | ~index. + + Flipping the sign bit (all bits when negative) maps IEEE order onto + unsigned order. The index rides inverted in the low bits so ties resolve to + the smaller index, matching torch.topk. NaN is forced above +inf because + torch.topk ranks it highest, and _block_scores_kernel can produce it. + Invalid lanes pack to 0, below every valid key. + """ + bits = score.to(tl.uint32, bitcast=True) + ordered = bits ^ tl.where(bits >> 31 != 0, 0xFFFFFFFF, 0x80000000) + ordered = tl.where(score != score, 0xFFFFFFFF, ordered) + tie = (_IDX_MASK - index).to(tl.int64) + return tl.where(valid, (ordered.to(tl.int64) << _IDX_BITS) | tie, 0) + + +@triton.jit +def _scan_range_topk(scores, row, lo, hi, nblocks, K_PAD: tl.constexpr, TILE): + """Running top-K of packed keys over blocks [lo, hi) of one row.""" + best = tl.zeros((K_PAD,), dtype=tl.int64) + for start in range(lo, hi, TILE): + cols = start + tl.arange(0, TILE) + valid = cols < hi + score = tl.load(scores + row * nblocks + cols, valid, other=_NEG_INF) + # -inf marks padding; it must never take a slot. + key = _pack_key(score, cols, valid & (score != _NEG_INF)) + best = tl.topk(tl.join(best, tl.topk(key, K_PAD)).reshape(2 * K_PAD), K_PAD) + return best + + +@triton.jit +def _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT): + """Blocks this row can populate, matching _block_scores_kernel's extent. + + The logits workspace is padded to the full width, so past a row's end every + block scores -inf and can never win a slot. Bounding the scan here is what + torch.topk cannot do: it has no per-row extent. + """ + bound = row // ROW_REPEAT + start = tl.load(starts + bound) if HAS_STARTS else 0 + end = tl.load(ends + bound) + live = (end - start + BLOCK_SIZE - 1) // BLOCK_SIZE + return tl.minimum(tl.maximum(live, 1), nblocks).to(tl.int32) + + +@triton.jit +def _active_chunks(live, NUM_CHUNKS: tl.constexpr, TILE: tl.constexpr): + """Chunks worth running: each should own a full tile, else short contexts + pay NUM_CHUNKS tile-wide top-ks to look at almost nothing.""" + return tl.minimum(tl.maximum((live + TILE - 1) // TILE, 1), NUM_CHUNKS) + + +@triton.jit +def _write_candidates( + best, + out, + row, + stride_row, + stride_col, + OUT_K: tl.constexpr, + K_PAD: tl.constexpr, + TAIL: tl.constexpr, +): + slots = tl.arange(0, K_PAD) + index = (_IDX_MASK - (best & _IDX_MASK)).to(tl.int32) + tl.store( + out + row * stride_row + slots * stride_col, + tl.where(best != 0, index, -1), + slots < OUT_K, + ) + if TAIL > 0: + # Fewer blocks than requested candidates: pad the unused slots. + rest = K_PAD + tl.arange(0, TAIL) + tl.store(out + row * stride_row + rest * stride_col, -1, rest < OUT_K) + + +@triton.jit(do_not_specialize=["nblocks"]) +def _row_topk_kernel( + scores, + starts, + ends, + out, + nblocks, + stride_row, + stride_col, + BLOCK_SIZE: tl.constexpr, + HAS_STARTS: tl.constexpr, + ROW_REPEAT: tl.constexpr, + OUT_K: tl.constexpr, + K_PAD: tl.constexpr, + TAIL: tl.constexpr, + TILE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + hi = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT) + best = _scan_range_topk(scores, row, 0, hi, nblocks, K_PAD, TILE) + _write_candidates(best, out, row, stride_row, stride_col, OUT_K, K_PAD, TAIL) + + +@triton.jit(do_not_specialize=["nblocks"]) +def _row_topk_partial_kernel( + scores, + starts, + ends, + partial, + nblocks, + BLOCK_SIZE: tl.constexpr, + HAS_STARTS: tl.constexpr, + ROW_REPEAT: tl.constexpr, + NUM_CHUNKS: tl.constexpr, + K_PAD: tl.constexpr, + TILE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + chunk = tl.program_id(1) + live = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT) + # Split on device from this row's live extent: a host-side bound would cost + # a sync, and the padded width says nothing about the context. + active = _active_chunks(live, NUM_CHUNKS, TILE) + if chunk < active: + split = (live + active - 1) // active + lo = chunk * split + hi = tl.minimum(lo + split, live) + best = _scan_range_topk(scores, row, lo, hi, nblocks, K_PAD, TILE) + slots = tl.arange(0, K_PAD) + tl.store(partial + row * (NUM_CHUNKS * K_PAD) + chunk * K_PAD + slots, best) + + +@triton.jit(do_not_specialize=["nblocks"]) +def _row_topk_merge_kernel( + partial, + starts, + ends, + out, + nblocks, + stride_row, + stride_col, + BLOCK_SIZE: tl.constexpr, + HAS_STARTS: tl.constexpr, + ROW_REPEAT: tl.constexpr, + NUM_CHUNKS: tl.constexpr, + OUT_K: tl.constexpr, + K_PAD: tl.constexpr, + TAIL: tl.constexpr, + TILE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + live = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT) + active = _active_chunks(live, NUM_CHUNKS, TILE) + lanes = tl.arange(0, NUM_CHUNKS * K_PAD) + # Chunks the partial pass skipped hold stale keys; mask them out. + keys = tl.load( + partial + row * (NUM_CHUNKS * K_PAD) + lanes, + (lanes // K_PAD) < active, + other=0, + ) + _write_candidates( + tl.topk(keys, K_PAD), out, row, stride_row, stride_col, OUT_K, K_PAD, TAIL + ) + + +# One wave of workgroups on MI355X; below this a program-per-row leaves most +# CUs idle and splitting the row pays for the extra merge. +_TARGET_PROGRAMS = 256 +# The merge top-k spans NUM_CHUNKS*K_PAD lanes whether or not those chunks ran, +# so more chunks buy parallelism at a cost the low-row cases cannot absorb. +_MAX_CHUNKS = 8 +_TILE = 512 + + +def _rocm_select_topk( + scores: torch.Tensor, + row_ks: torch.Tensor | None, + row_ke: torch.Tensor, + nblocks: int, + topk_blocks: int, + out: torch.Tensor, + row_repeat: int, + block_size: int, +) -> None: + """Fused block top-k for ROCm, replacing topk + _store_candidates_kernel. + + torch.topk runs a multi-pass radix select that sorts values it does not + need -- the caller only consumes the *set* of block ids -- and scans the + padding because it cannot see per-row extents. + """ + rows = scores.shape[0] + k_pad = triton.next_power_of_2(min(topk_blocks, nblocks)) + # _scan_range_topk does a running top-K_PAD merge one TILE-wide slice at a + # time; tl.topk silently corrupts its output if asked for more elements + # than the slice holds, so the tile must never be smaller than K_PAD. + tile = max(_TILE, k_pad) + + num_chunks = 1 + if rows < _TARGET_PROGRAMS: + # Sized only to fill the machine; how many actually run is decided per + # row on device from its live extent. + want = min( + triton.cdiv(_TARGET_PROGRAMS, max(rows, 1)), + _MAX_CHUNKS, + max(1, nblocks // tile), + ) + num_chunks = max(1, triton.next_power_of_2(want + 1) // 2) + + common = dict( + BLOCK_SIZE=block_size, + HAS_STARTS=row_ks is not None, + ROW_REPEAT=row_repeat, + OUT_K=topk_blocks, + K_PAD=k_pad, + TAIL=triton.next_power_of_2(topk_blocks - k_pad) if topk_blocks > k_pad else 0, + TILE=tile, + ) + if num_chunks == 1: + _row_topk_kernel[(rows,)]( + scores, row_ks, row_ke, out, nblocks, *out.stride(), **common + ) + return + + partial = scores.new_empty((rows, num_chunks * k_pad), dtype=torch.int64) + _row_topk_partial_kernel[(rows, num_chunks)]( + scores, + row_ks, + row_ke, + partial, + nblocks, + BLOCK_SIZE=block_size, + HAS_STARTS=row_ks is not None, + ROW_REPEAT=row_repeat, + NUM_CHUNKS=num_chunks, + K_PAD=k_pad, + TILE=tile, + ) + _row_topk_merge_kernel[(rows,)]( + partial, + row_ks, + row_ke, + out, + nblocks, + *out.stride(), + NUM_CHUNKS=num_chunks, + **common, + ) + + def select_candidate_blocks( logits: torch.Tensor, row_ks: torch.Tensor | None, @@ -176,6 +427,11 @@ def select_candidate_blocks( row_repeat, 128, ) + if current_platform.is_rocm(): + _rocm_select_topk( + scores, row_ks, row_ke, nblocks, topk_blocks, out, row_repeat, block_size + ) + 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))](