[ROCm][Perf] Replace torch.topk in DSA candidate block selection

[Fangzhou-Ai]

Summary

select_candidate_blocks picks the top-k blocks per row with torch.topk, which runs a multi-pass radix select (computeBlockDigitCounts, computeBlockwiseWithinKCounts, gatherTopK). That does two things this caller does not need:

• It orders the results. The only consumer is _candidate_flags_kernel, which does flags[block] = 1. Nothing downstream reads the ranking or the scores — only the set of selected block ids matters.
• It scans the whole padded workspace. The logits workspace is padded to max_model_len, so nblocks is max_model_len / block_size no matter how short the context is. Past a row's end every block scores -inf and can never win a slot, but torch.topk has no way to express a per-row extent, so it reads all of it.

This PR replaces that select on ROCm with a Triton kernel:

• Each score and its index are packed into one order-preserving int64 key, so tl.topk returns both at once and the separate _store_candidates_kernel pass disappears. The float bits are sign-transformed so IEEE order becomes unsigned order, and the index rides inverted in the low bits so ties resolve to the smaller index, matching torch.topk.
• Each row is bounded by its own end, which is the part torch.topk cannot do. At 4K context in a 256K workspace this skips ~98% of the columns.
• Decode rows are batch * next_n, only 1–640, which is too few to fill the GPU one program per row. Rows are split across programs and merged, with the split size computed on device from each row's live extent — a host-side bound would need a sync, and sizing from the padded width instead measured 6.3x → 2.0x at low concurrency.

The CUDA path is unchanged; the new kernel is behind current_platform.is_rocm().

Not a duplicate

Checked open PRs touching this area:

• [ROCm][Perf] Candidate-only MQA scorer for DSA sparse indexer prefill on gfx950 #57858 is the only other open PR touching candidate_blocks.py, but it optimizes prefill consumer-layer scoring (candidate_mqa_logits, scoring only candidate columns). The producer-side block select this PR rewrites is untouched there. The two are complementary. Note [ROCm][Perf] Candidate-only MQA scorer for DSA sparse indexer prefill on gfx950 #57858 also flips _store_candidates_kernel's NaN handling (!= → >); the ROCm path here bypasses that kernel entirely, so the behaviors do not collide.
• [Perf][Kernel] Fused DSA indexer Top-k kernel (LiteTopk) #48726 (LiteTopk) is the indexer's token top-k, a different kernel and stage.
• [ROCm] Route sparse-indexer decode top-k through AITER dispatcher #50470, [ROCm] [Perf] Optimize rocm aiter sparse mla indexer #46172, [ROCm][Perf] Shard DSA prefill indexer rows across TP ranks #56944 touch rocm_aiter_mla_sparse.py only (AITER dispatch, indexer optimization, prefill row sharding), not block candidate selection.

No open PR optimizes select_candidate_blocks.

Correctness

The new path must select the same block set as the existing path. Verified against the unchanged path across 27 configurations — contexts 512 to 262144, rows 1 to 640, ragged per-row ends, non-zero row_ks starts, row_repeat decode grouping, topk_blocks > nblocks, and NaN scores:

27 cases, 0 mismatches

Two cases this caught during development, both now handled:

• torch.topk ranks NaN highest, and _block_scores_kernel can emit NaN during warmup, so keys force NaN above +inf.
• When nblocks < topk_blocks the internal K_PAD shrinks below the output width, leaving tail columns unwritten; they are now -1 filled.

Tests

$ pytest tests/kernels/attention/test_rocm_aiter_candidate_mask.py -q
9 passed, 14 warnings in 4.23s

$ pre-commit run --files vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
exit=0 (no failures)

Kernel performance (MI355X)

Full select_candidate_blocks, CUDA-graph timed, swept over the decode shape space: concurrency 1–128 x spec-decode depth 1–5, max_model_len=262144, block_size=8, k=64. All 120 points are faster.

context	median	range
4096	4.16x	3.97–5.08x
16384	3.99x	3.27–4.70x
65536	2.94x	1.64–3.15x

Representative points (microseconds):

context	conc	next_n	rows	before	after	speedup
4096	1	1	1	47.55	11.20	4.24x
4096	32	1	32	58.00	13.96	4.16x
4096	128	5	640	293.54	67.86	4.33x
16384	32	1	32	59.19	14.31	4.14x
16384	128	5	640	299.65	91.67	3.27x
65536	32	1	32	60.31	19.30	3.13x
65536	128	5	640	330.51	185.83	1.78x

These shapes are TP-independent: the DSA indexer is replicated across ranks (index_n_heads is not divided by tp_size, unlike the main attention heads), so TP2 and TP4 run identical selections and one sweep covers both.

Notes

The kernel is still well above the bandwidth roofline (~45x at the largest shapes, more at small ones where it is latency-bound on two launches), so there is headroom left. TILE=512 is a compromise: 1024 would take 65536-context from 2.36x to 2.79x but drop 4096-context from 23.8x to 14.4x. A context-aware tile choice could get both.

Benchmarks used uniform context per batch; heavily ragged batches are bounded by their longest row.

End-to-end serving

Measured on MI355X TP4, DeepSeek-V4.1-Flash, dspark MTP (5 speculative tokens, acceptance pinned at 3.51 so acceptance variance cannot contribute). Both arms are the same build differing only by an env toggle on the selection path, run concurrently on separate 4-GPU sets, after a warmup pass over every shape so Triton JIT and graph capture are paid before measurement.

snapshot of trace:
Before:

After:

There is no measurable end-to-end difference, and none should be expected. Output throughput, ISL 4096 / OSL 128:

concurrency	baseline tok/s	this PR tok/s
1	259.4	257.1
8	1287.7	1282.4
32	1068.9	1070.3
128	1253.8	1254.8

Repeating a single configuration three times shows run-to-run spread of 1112–1191 tok/s on baseline and 1142–1166 on this branch, with the sign of the difference flipping between repeats. The noise floor is several percent, while the expected effect is about 0.2%: the selection saves roughly 45 us per step against a 14.6 ms inter-token latency. The kernel win is real but it is a sub-1% slice of step time, so it cannot show up above noise in serving throughput.

Median ITL, the most stable metric here, is unchanged (14.59-14.64 ms baseline vs 14.61-14.73 ms this branch).

This PR should therefore be read as removing a hotspot that is prominent in kernel traces and cutting its cost 2-5x, not as an end-to-end throughput win.

AI assistance

AI assistance was used for this change.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Fangzhou-Ai]

Hi @tjtanaa can you take a look at this new kernel?

[github-actions]

✅ @Fangzhou-Ai, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91382 for commit 7f00fc957d67.

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91450 for commit 410277d72976.

[Fangzhou-Ai]

/ci run

[github-actions]

❌ This PR is 19 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91525 for commit 1f07edd53272.