[Perf][Kernel] Add sampled filtering for persistent top-k

[mgoin]

Purpose

Speed up long FP32 sparse-indexer decode top-k with sampling inspired by DeepSelect. A coalesced sample estimates a cutoff, survivors are compacted in shared memory, and exact FP32 selection finishes on those candidates. Too few or too many candidates trigger exact full-row selection.

Enable sampling above 64 rows, with actual valid length ≥98,304 for k=512 or ≥65,536 for k=1024/2048, and at least 144 KiB opt-in shared memory. Shorter runtime rows reuse the existing selector within the same launch. These conservative bounds come from the fine B300 sweep below.

The plot forces sampling to expose the crossover; it shows the minimum speedup across measured batches, allocations, seeds and cache modes. The k=512 slowdown near 80K–86K motivates its later cutoff.

Final dispatch improves every tested case entirely within the sampling region. At B=1024 with 80–100% ragged lengths, warm latency drops 316.58→207.06 µs at 128K/k=2048 and 584.45→301.72 µs at 256K/k=512. Short-row fallback can add up to 3.32%; the largest absolute increase measured was 1.215 µs. Performance validation is B300-only.

Duplicate checks found no open PR implementing this sampled filter: #44606 tunes existing dispatch, #55872 adds a FlashInfer backend, and #55314 fixes existing stash overflow. That overflow remains outside the new path. Prefill is unchanged.

Test Plan

Added graph-replay correctness coverage and benchmarks/kernels/benchmark_persistent_topk.py for full/ragged lengths and warm/cold L2. All GPU runs used chg reservations and source-built baseline/candidate libraries.

Local validation commands run

The wrappers load the isolated candidate library and invoke tests/kernels/test_top_k_per_row.py.

chg run -- .venv/bin/python benchmarks/results/topk_sweep_20260910/run_tests.py
chg run -- .venv/bin/python benchmarks/results/topk_sweep_20260910/run_tests96.py \
  -k "sampled_graph or workspace_topk or reused_group"
chg run -- .venv/bin/python benchmarks/results/topk_sweep_20260910/run_tests96.py -k sampled_graph
chg run -- compute-sanitizer --tool memcheck --error-exitcode 1 \
  .venv/bin/python benchmarks/results/topk_sweep_20260910/sanitize96.py
chg run -- compute-sanitizer --tool synccheck --error-exitcode 1 \
  .venv/bin/python benchmarks/results/topk_sweep_20260910/sanitize96.py --fallback
.venv/bin/pre-commit run --files csrc/libtorch_stable/{sampled_topk.cuh,persistent_topk.cuh,topk.cu} \
  benchmarks/kernels/benchmark_persistent_topk.py tests/kernels/test_top_k_per_row.py

Test Result

• 6,265 forced-backend cases and 541 final-dispatch cases: zero selection errors.
• Full suite after helper extraction: 201 passed, 19 skipped. Final cutoff checks: 125 passed, 6 skipped; all 30 graph cases passed again after the shared-memory test guard.
• Memcheck/synccheck: zero errors. Pre-commit passed; SM103a and SM80 compile.
• DeepSeek-V4-Flash-0731, TP2, 128 requests with a 401,408-token shared prefix: both arms passed exact FP32 top-k audits at 100K valid scores and answered 128/128 arithmetic questions correctly. Both scored 123/128 on the GSM8K subset, which used the unchanged cooperative path. Long generations matched through EOS for 75/128 requests; this does not establish token equivalence or serving speedup. The local model harness used source-built top-k libraries and an adapter for inactive optional MoE arguments in the older installed extension.

AI assistance was used for implementation, benchmarking, 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.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88352 for commit 93309ace0e7e.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88691 for commit fa8d8fdf62db.