[DarkSharpness]

(auto generated by Claude)

Speeds up DeepSelect top-k for small-batch, long-context decode (the DeepSeek-V4.1 indexer case) and lets the V4.1 index layers run on bf16 scores.

DeepSelect cluster kernel (bf16, topk <= 1024)

• Rewritten cluster kernel: each CTA takes a contiguous slice of the row; a slice of up to 32K stays in registers, a longer one streams through a bulk-copy ring behind a cluster-wide pivot taken from a window (three pseudo-random stages plus the slice's tail). The cluster agrees on the final pivot by pushing 256-bin histograms to every CTA, then writes its output directly (DSMEM atomics for unsorted, a bitmap prefix for sorted_index).
• Dispatch keys only on the row count, never on the (CUDA-graph padded) width: cluster sizes 16/8/4/3/2 are tiered by an occupancy probe of how many clusters fit in one wave; the per-row end decides everything else inside the kernel.
• All modules of a call signature are compiled in parallel (parallel_compile in the JIT utils).
• Without abort, a NaN row now selects index 0 in every slot instead of writing a 0x3F3F3F3F sentinel to slot 0; return_value and abort_when_nan_found default to False.

DeepSeek-V4.1 indexer

• SGLANG_OPT_DSV41_BF16_INDEX_LOGITS (default off): the DeepGEMM fp4 MQA logits are emitted in bf16 (the reference ranks bf16 scores), their top-k goes through DeepSelect, and amax8_varlen reads bf16 scores (now templated on input/output dtype).

Kernel latency (B200, cold L2, width 1M, topk 1024, us at end 128K / 1M). Before, the cluster kernel only ran for rows <= 6 at width >= 512K; every other case took one CTA per row.

rows	one CTA per row	cluster (this PR)
1	19.0 / 81.8	6.1 / 13.0
8	19.3 / 82.7	6.5 / 15.5
32	19.5 / 82.8	8.6 / 22.3
66	19.8 / 83.2	13.2 / 40.9

Validation

• B200 and H200: test_deep_select.py, plus randomized checks of every cluster size against torch.topk (ties, constants, -inf, +-0, NaN, ascending/descending rows, ragged end, offsets, int64, page transform, values, sorted_index).
• B200: DeepSeek-V4.1 indexer kernel tests pass with the flag off; with it on, only tests whose references use fp32 logits differ (bf16 rounding and ties at the selection floor).
• 2x GB300, DeepSeek-V4.1-Flash end to end, flag on vs off: no crash, GSM8K within run-to-run noise (0.893 vs 0.913 / 0.897 on a rerun), passkey correct at 125K and 209K; decode top-k in the index layers 51.1 -> 27.6 us per step at 128K.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): 🚫 Run #38555200909
Latest PR Test (Extra): ❌ Run #38555200886
Latest PR Test (AMD ROCm 10): ❌ Run #38555200764

[yiakwy-xpu-ml-framework-team]

@DarkSharpness we are supporting epic long context with our brand new sparse attention , where top 2048 is very useful !

Let me explain, we will 2048 chunks, with partials scores each chunk corresponding to 64~128 elements, where we can select total 1 billion context efficiently.