Skip to content

[Bugfix] Revert "[ROCm][Perf] Replace torch.topk in DSA candidate block selection" #59125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ywang96 merged 1 commit into from
Sep 29, 2026
Merged

[Bugfix] Revert "[ROCm][Perf] Replace torch.topk in DSA candidate block selection" #59125

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
256 changes: 0 additions & 256 deletions vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,8 @@

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):
Expand Down Expand Up @@ -143,252 +138,6 @@ 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,
Expand Down Expand Up @@ -427,11 +176,6 @@ 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))](
Expand Down
Loading