@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6b24bd8e75

[ROCm][Perf] Replace torch.topk in DSA candidate block selection (#58208)

Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>

Fangzhou Ai committed Sep 28, 2026 at 10:42 UTC 6b24bd8e75a3e465083e8a131968915bc00efaf4
1 file changed +256
vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
+256
@@ -3,8 +3,13 @@
3
4 import torch
5
6 +from vllm.platforms import current_platform
7 from vllm.triton_utils import tl, triton
8
9 +_IDX_BITS = tl.constexpr(20)
10 +_IDX_MASK = tl.constexpr((1 << 20) - 1)
11 +_NEG_INF = tl.constexpr(-float("inf"))
12 +
13
14 @triton.jit
15 def _max_with_nan(a, b):
@@ -138,6 +143,252 @@ def _mask_candidates_kernel(
143 )
144
145
146 +@triton.jit
147 +def _pack_key(score, index, valid):
148 + """Order-preserving int64 key: float order << _IDX_BITS | ~index.
149 +
150 + Flipping the sign bit (all bits when negative) maps IEEE order onto
151 + unsigned order. The index rides inverted in the low bits so ties resolve to
152 + the smaller index, matching torch.topk. NaN is forced above +inf because
153 + torch.topk ranks it highest, and _block_scores_kernel can produce it.
154 + Invalid lanes pack to 0, below every valid key.
155 + """
156 + bits = score.to(tl.uint32, bitcast=True)
157 + ordered = bits ^ tl.where(bits >> 31 != 0, 0xFFFFFFFF, 0x80000000)
158 + ordered = tl.where(score != score, 0xFFFFFFFF, ordered)
159 + tie = (_IDX_MASK - index).to(tl.int64)
160 + return tl.where(valid, (ordered.to(tl.int64) << _IDX_BITS) | tie, 0)
161 +
162 +
163 +@triton.jit
164 +def _scan_range_topk(scores, row, lo, hi, nblocks, K_PAD: tl.constexpr, TILE):
165 + """Running top-K of packed keys over blocks [lo, hi) of one row."""
166 + best = tl.zeros((K_PAD,), dtype=tl.int64)
167 + for start in range(lo, hi, TILE):
168 + cols = start + tl.arange(0, TILE)
169 + valid = cols < hi
170 + score = tl.load(scores + row * nblocks + cols, valid, other=_NEG_INF)
171 + # -inf marks padding; it must never take a slot.
172 + key = _pack_key(score, cols, valid & (score != _NEG_INF))
173 + best = tl.topk(tl.join(best, tl.topk(key, K_PAD)).reshape(2 * K_PAD), K_PAD)
174 + return best
175 +
176 +
177 +@triton.jit
178 +def _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT):
179 + """Blocks this row can populate, matching _block_scores_kernel's extent.
180 +
181 + The logits workspace is padded to the full width, so past a row's end every
182 + block scores -inf and can never win a slot. Bounding the scan here is what
183 + torch.topk cannot do: it has no per-row extent.
184 + """
185 + bound = row // ROW_REPEAT
186 + start = tl.load(starts + bound) if HAS_STARTS else 0
187 + end = tl.load(ends + bound)
188 + live = (end - start + BLOCK_SIZE - 1) // BLOCK_SIZE
189 + return tl.minimum(tl.maximum(live, 1), nblocks).to(tl.int32)
190 +
191 +
192 +@triton.jit
193 +def _active_chunks(live, NUM_CHUNKS: tl.constexpr, TILE: tl.constexpr):
194 + """Chunks worth running: each should own a full tile, else short contexts
195 + pay NUM_CHUNKS tile-wide top-ks to look at almost nothing."""
196 + return tl.minimum(tl.maximum((live + TILE - 1) // TILE, 1), NUM_CHUNKS)
197 +
198 +
199 +@triton.jit
200 +def _write_candidates(
201 + best,
202 + out,
203 + row,
204 + stride_row,
205 + stride_col,
206 + OUT_K: tl.constexpr,
207 + K_PAD: tl.constexpr,
208 + TAIL: tl.constexpr,
209 +):
210 + slots = tl.arange(0, K_PAD)
211 + index = (_IDX_MASK - (best & _IDX_MASK)).to(tl.int32)
212 + tl.store(
213 + out + row * stride_row + slots * stride_col,
214 + tl.where(best != 0, index, -1),
215 + slots < OUT_K,
216 + )
217 + if TAIL > 0:
218 + # Fewer blocks than requested candidates: pad the unused slots.
219 + rest = K_PAD + tl.arange(0, TAIL)
220 + tl.store(out + row * stride_row + rest * stride_col, -1, rest < OUT_K)
221 +
222 +
223 +@triton.jit(do_not_specialize=["nblocks"])
224 +def _row_topk_kernel(
225 + scores,
226 + starts,
227 + ends,
228 + out,
229 + nblocks,
230 + stride_row,
231 + stride_col,
232 + BLOCK_SIZE: tl.constexpr,
233 + HAS_STARTS: tl.constexpr,
234 + ROW_REPEAT: tl.constexpr,
235 + OUT_K: tl.constexpr,
236 + K_PAD: tl.constexpr,
237 + TAIL: tl.constexpr,
238 + TILE: tl.constexpr,
239 +):
240 + row = tl.program_id(0).to(tl.int64)
241 + hi = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT)
242 + best = _scan_range_topk(scores, row, 0, hi, nblocks, K_PAD, TILE)
243 + _write_candidates(best, out, row, stride_row, stride_col, OUT_K, K_PAD, TAIL)
244 +
245 +
246 +@triton.jit(do_not_specialize=["nblocks"])
247 +def _row_topk_partial_kernel(
248 + scores,
249 + starts,
250 + ends,
251 + partial,
252 + nblocks,
253 + BLOCK_SIZE: tl.constexpr,
254 + HAS_STARTS: tl.constexpr,
255 + ROW_REPEAT: tl.constexpr,
256 + NUM_CHUNKS: tl.constexpr,
257 + K_PAD: tl.constexpr,
258 + TILE: tl.constexpr,
259 +):
260 + row = tl.program_id(0).to(tl.int64)
261 + chunk = tl.program_id(1)
262 + live = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT)
263 + # Split on device from this row's live extent: a host-side bound would cost
264 + # a sync, and the padded width says nothing about the context.
265 + active = _active_chunks(live, NUM_CHUNKS, TILE)
266 + if chunk < active:
267 + split = (live + active - 1) // active
268 + lo = chunk * split
269 + hi = tl.minimum(lo + split, live)
270 + best = _scan_range_topk(scores, row, lo, hi, nblocks, K_PAD, TILE)
271 + slots = tl.arange(0, K_PAD)
272 + tl.store(partial + row * (NUM_CHUNKS * K_PAD) + chunk * K_PAD + slots, best)
273 +
274 +
275 +@triton.jit(do_not_specialize=["nblocks"])
276 +def _row_topk_merge_kernel(
277 + partial,
278 + starts,
279 + ends,
280 + out,
281 + nblocks,
282 + stride_row,
283 + stride_col,
284 + BLOCK_SIZE: tl.constexpr,
285 + HAS_STARTS: tl.constexpr,
286 + ROW_REPEAT: tl.constexpr,
287 + NUM_CHUNKS: tl.constexpr,
288 + OUT_K: tl.constexpr,
289 + K_PAD: tl.constexpr,
290 + TAIL: tl.constexpr,
291 + TILE: tl.constexpr,
292 +):
293 + row = tl.program_id(0).to(tl.int64)
294 + live = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT)
295 + active = _active_chunks(live, NUM_CHUNKS, TILE)
296 + lanes = tl.arange(0, NUM_CHUNKS * K_PAD)
297 + # Chunks the partial pass skipped hold stale keys; mask them out.
298 + keys = tl.load(
299 + partial + row * (NUM_CHUNKS * K_PAD) + lanes,
300 + (lanes // K_PAD) < active,
301 + other=0,
302 + )
303 + _write_candidates(
304 + tl.topk(keys, K_PAD), out, row, stride_row, stride_col, OUT_K, K_PAD, TAIL
305 + )
306 +
307 +
308 +# One wave of workgroups on MI355X; below this a program-per-row leaves most
309 +# CUs idle and splitting the row pays for the extra merge.
310 +_TARGET_PROGRAMS = 256
311 +# The merge top-k spans NUM_CHUNKS*K_PAD lanes whether or not those chunks ran,
312 +# so more chunks buy parallelism at a cost the low-row cases cannot absorb.
313 +_MAX_CHUNKS = 8
314 +_TILE = 512
315 +
316 +
317 +def _rocm_select_topk(
318 + scores: torch.Tensor,
319 + row_ks: torch.Tensor | None,
320 + row_ke: torch.Tensor,
321 + nblocks: int,
322 + topk_blocks: int,
323 + out: torch.Tensor,
324 + row_repeat: int,
325 + block_size: int,
326 +) -> None:
327 + """Fused block top-k for ROCm, replacing topk + _store_candidates_kernel.
328 +
329 + torch.topk runs a multi-pass radix select that sorts values it does not
330 + need -- the caller only consumes the *set* of block ids -- and scans the
331 + padding because it cannot see per-row extents.
332 + """
333 + rows = scores.shape[0]
334 + k_pad = triton.next_power_of_2(min(topk_blocks, nblocks))
335 + # _scan_range_topk does a running top-K_PAD merge one TILE-wide slice at a
336 + # time; tl.topk silently corrupts its output if asked for more elements
337 + # than the slice holds, so the tile must never be smaller than K_PAD.
338 + tile = max(_TILE, k_pad)
339 +
340 + num_chunks = 1
341 + if rows < _TARGET_PROGRAMS:
342 + # Sized only to fill the machine; how many actually run is decided per
343 + # row on device from its live extent.
344 + want = min(
345 + triton.cdiv(_TARGET_PROGRAMS, max(rows, 1)),
346 + _MAX_CHUNKS,
347 + max(1, nblocks // tile),
348 + )
349 + num_chunks = max(1, triton.next_power_of_2(want + 1) // 2)
350 +
351 + common = dict(
352 + BLOCK_SIZE=block_size,
353 + HAS_STARTS=row_ks is not None,
354 + ROW_REPEAT=row_repeat,
355 + OUT_K=topk_blocks,
356 + K_PAD=k_pad,
357 + TAIL=triton.next_power_of_2(topk_blocks - k_pad) if topk_blocks > k_pad else 0,
358 + TILE=tile,
359 + )
360 + if num_chunks == 1:
361 + _row_topk_kernel[(rows,)](
362 + scores, row_ks, row_ke, out, nblocks, *out.stride(), **common
363 + )
364 + return
365 +
366 + partial = scores.new_empty((rows, num_chunks * k_pad), dtype=torch.int64)
367 + _row_topk_partial_kernel[(rows, num_chunks)](
368 + scores,
369 + row_ks,
370 + row_ke,
371 + partial,
372 + nblocks,
373 + BLOCK_SIZE=block_size,
374 + HAS_STARTS=row_ks is not None,
375 + ROW_REPEAT=row_repeat,
376 + NUM_CHUNKS=num_chunks,
377 + K_PAD=k_pad,
378 + TILE=tile,
379 + )
380 + _row_topk_merge_kernel[(rows,)](
381 + partial,
382 + row_ks,
383 + row_ke,
384 + out,
385 + nblocks,
386 + *out.stride(),
387 + NUM_CHUNKS=num_chunks,
388 + **common,
389 + )
390 +
391 +
392 def select_candidate_blocks(
393 logits: torch.Tensor,
394 row_ks: torch.Tensor | None,
@@ -176,6 +427,11 @@ def select_candidate_blocks(
427 row_repeat,
428 128,
429 )
430 + if current_platform.is_rocm():
431 + _rocm_select_topk(
432 + scores, row_ks, row_ke, nblocks, topk_blocks, out, row_repeat, block_size
433 + )
434 + return
435 # Keep the existing top-k tie behavior.
436 top = scores.topk(min(topk_blocks, nblocks), dim=-1)
437 _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](