[Kernel] Optimize CuTeDSL DCP top-k merge

[LucasWilkinson]

Summary

Two commits optimizing the DCP sparse-indexer top-k merge (each rank holds the local top-K of its 1/DCP KV shard; the merge selects the global top-K from the gathered candidates).

Commit 1 — CuTeDSL selector optimization. StableTopKFromGatheredCandidatesKernel runs a radix pass over 32-bit score keys first and only runs the token-id tie-break pass when the score threshold actually ties, instead of always sorting 64-bit score<<32 | ~id keys. Selection semantics are unchanged and deterministic: score descending, then lowest global token id.

Commit 2 — symmetric-memory candidate exchange. Replaces the NCCL all-gather of packed (score, global_id) candidates with a one-shot all-gather over torch.distributed._symmetric_memory: each rank packs its candidates into its slice of every peer's symmetric "inbox" with a Triton remote-write kernel, synchronized by device-side write/read sequence flags (sys-scope acquire/release atomics). No NCCL launch, no host branches, fully CUDA-graph capturable; the selector then runs on the local inbox unchanged. get_dcp_topk_symm_mem_workspace returns None on unsupported configurations (no device group, unaligned candidate counts, rendezvous failure) and the code falls back to the existing all-gather path.

Duplicate-work check

gh pr list --repo vllm-project/vllm --state open --search "CuTeDSL DCP top-k merge" --limit 20
gh pr list --repo vllm-project/vllm --state open --search "sparse indexer topk cutedsl" --limit 20
gh pr list --repo vllm-project/vllm --state open --search "DCP topk symmetric memory" --limit 20

Related DCP/sparse-MLA PRs (#46514, #45426) exist but none touch the CuTeDSL gathered-candidate selector or the candidate exchange transport. The side-stream indexer overlap work is separate in #47355.

Benchmarks

Selector microbenchmark (stable_topk_from_gathered_candidates_cutedsl, median µs, 7×200 launches, JIT warmup excluded):

case	main	this PR	speedup
rows=16, candidates=4096, random	6.922	6.721	1.03x
rows=16, candidates=8192, random	13.040	11.478	1.14x
rows=16, candidates=16384, random	20.936	20.684	1.01x
rows=64, candidates=8192, random	13.029	11.696	1.11x
rows=256, candidates=8192, random	17.705	15.459	1.15x
rows=16, candidates=8192, tie-heavy	21.371	20.396	1.05x

The symm-mem exchange replaces an NCCL all-gather launch + copy with two flag kernels and one push kernel on the capture stream; on GLM-5.1-NVFP4 B200 TP=4/DCP=4 decode it was worth ~4% end-to-end output throughput measured back-to-back against the NCCL merge path.

Tests

.venv/bin/python -m pytest tests/v1/attention/test_indexer_dcp_localize.py -q       # 38 passed
CUDA_VISIBLE_DEVICES=0,1 .venv/bin/python -m pytest tests/distributed/test_dcp_topk_symm_mem.py -q   # passed (2- and 4-rank)
pre-commit run --files <changed files>   # all hooks passed

tests/distributed/test_dcp_topk_symm_mem.py is new: multi-rank exactness of the symm-mem merge against the all-gather reference (including tie-heavy scores and empty shards) and replay correctness under CUDA graph capture.

AI assistance

AI assistance (Claude) was used to implement, benchmark, and prepare this PR. The human submitter has reviewed the changed lines and run the tests above.

[Leoyzen]

Missing sync_threads() between score_threshold = prefix_smem.load() above and this store(0) — warps that haven't loaded yet capture score_threshold = 0, so the phase-2 key == score_threshold predicate drops every tied candidate. Adding one barrier before this store makes 20/20 heavy-tie trials match the 64-bit reference bit-exactly.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork