[ROCm][Perf] Optimize DSV4.1 K=512 decode top-k on gfx950

[Fangzhou-Ai]

Purpose

Part of #56506.

DeepSeek-V4.1-Flash uses K=512 for decode indexer top-k, which misses the
existing gfx950 tuning and falls back to the generic 10-split kernel. Add a
ROCm/gfx950-only K=512 path for up to 384 rows and 1M columns. It selects
active splits from device sequence lengths so FULL graph replay remains valid,
reduces masked -inf histogram contention, and uses a smaller final sort when
possible. Other architectures, K values, strides, and shapes keep the existing
path.

This PR contains no multistream CSA, scheduler, or queue-management changes.

Results

• All 1,008 matrix cases passed: concurrency 1/2/4/8/16/32/64,
  speculative tokens 1-6, context 10K/50K/100K/200K/500K/1M, compression
  ratios 1/2, and dense/masked logits.

Cache condition	Cases	Geomean	Median	Speedup range	Slower cells
Warm HIP graph	1,008	2.310x	2.191x	1.008-12.344x	0
256 MiB cache pressure	1,008	1.824x	1.730x	1.023-6.246x	0

Representative CR1 masked decode results from the warm matrix:

Concurrency	Spec tokens	Context	Baseline (us)	Candidate (us)	Speedup
4	6	10,000	31.56	10.10	3.124x
4	6	50,000	34.79	15.80	2.202x
4	6	100,000	40.27	18.63	2.162x
4	6	200,000	51.78	21.08	2.456x
4	6	500,000	90.72	35.47	2.558x
4	6	1,000,000	148.10	49.31	3.003x
16	6	10,000	49.11	10.14	4.841x
16	6	50,000	60.68	21.98	2.761x
16	6	100,000	83.10	30.15	2.757x
16	6	200,000	121.48	41.79	2.907x
16	6	500,000	232.66	72.38	3.215x
16	6	1,000,000	420.57	135.95	3.094x

• A GPU trace captured 2,520 kernel intervals; 72 profiled pairs had a
  2.291x geometric-mean GPU busy-time speedup.
• Paired GSM8K strict accuracy was 1279/1319 baseline and 1286/1319
  candidate, with no invalid or truncated answers. Both arms passed 10K and
  100K token retrieval checks.
• Two short serving repetitions did not show a clear end-to-end throughput
  change, so this PR makes no serving-speed claim.

Validation

cmake --build dsv41_rocm_bench/topk/pr-cmake-build \
  --target _C_stable_libtorch -j16

HIP_VISIBLE_DEVICES=3 GPU_MAX_HW_QUEUES=4 .venv/bin/python -m pytest \
  tests/kernels/test_top_k_per_row.py \
  -k 'gfx950_long_c4a or gfx950_k512 or aiter_c4a' -v

.venv/bin/pre-commit run --files \
  csrc/libtorch_stable/sampler.cu \
  vllm/v1/attention/ops/rocm_aiter_mla_sparse.py \
  tests/kernels/test_top_k_per_row.py

The build passed, all 27 focused GPU tests passed, and scoped pre-commit passed.

This does not duplicate #54682, which is a MiniMax-M3 Triton selector, #50470,
which dispatches a separate sparse-indexer path to AITER, or #56628, which
optimizes the adjacent candidate mask.

OpenAI Codex assisted with implementation, benchmarks, tests, and analysis.

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

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88739 for commit 650bd5b911d9.