Add native SM120 and SM121 TopK with PyTorch stable ABI

[lucifer1004]

Summary

Add native SM120 and SM121 TopK support to the vLLM DeepSelect fork's dev branch, preserving its PyTorch stable ABI and existing SM100/SM103 backend. This is a DeepSelect library PR, not the vLLM engine integration.

• Native BF16/FP32 selection with int32/int64 indices, optional values, sorted indices or supported value sorting, variable row lengths and offsets.
• Exact native SM12 radix selection, streaming K512 selection, and segmented selection for small batches with long rows; no Torch selection fallback.
• One extension and one operator registration, retaining torch.ops.deep_select.topk, PyTorch 2.10 stable ABI, CPython ABI3, and TORCH_EXTENSION_NAME-derived initialization for downstream consumers.
• Architecture-isolated CUDA compilation: existing cluster/high-shared-memory kernels are not compiled for SM120/SM121. Explicit 12.0/12.1 and mixed builds, with real optional PTX emission. Default targets remain SM100a/SM103a.
• C++ device/current-stream, padding, storage-overlap, output-stride and scalar-bound checks also apply to direct operator calls. Nonempty-batch zero-width input is explicitly rejected on SM100/103, consistent with its original positive-width support; SM12 retains zero-width support.
• Preserve the target's optimized NaN handling, add bounded output stores, and retain the mandatory register-spill checker. Disable CUB outer-scan memoization in the native generic sorter to remove measured spills without changing ordering or precision.

Related upstream fix

This PR carries the cluster gather completion-barrier fix from deepseek-ai/DeepSelect#16, which is not yet in the current dev base. It maps both gather destinations and completion barriers to CTA0. The native SM12 backend does not use this cluster path. The upstream PR is independent and is not closed by this cross-reference.

Validation

Standard wheel builds use the existing artifact checker without disabling or weakening its thresholds. Tested builds include SM120-only, SM121-only, SM12 mixed, default SM100/103, and full SM100/103/120/121 mixed. The final full mixed build contains 348 kernel instances, all with zero STACK and LOCAL reported by the checker. An earlier source snapshot was also independently rebuilt from its sdist; that packaging test predates the final cluster-address and zero-width-contract corrections and is not presented as a final-snapshot sdist rebuild.

The wheel remains a single cp310-abi3 extension, with no unexpected undefined ATen/c10/torch C++ imports. Validation loads the actual extracted wheel rather than a source-tree shadow.

Final mixed-wheel validation (CUDA toolkit 13.2, driver 595.58.03, PyTorch 2.13.0+cu130, Compute Sanitizer 2026.1.1.0):

• SM120: 2,763 passed / 198 expected skips out of 2,961 cases, including 104 isolated device-trap cases and 12 cross-device/device-guard cases. Both the Python wrapper and direct registered operator were exercised. A risk-based selection of 19 streaming, segmented, generic, stream/graph and boundary cases passed memcheck and racecheck with zero errors/hazards/warnings; this is not full-suite sanitizer coverage.
• SM103/B300: 393 shared-contract tests passed / 4 multi-device cases skipped, including all 16 explicit zero-width rejections and 32 optimized NaN cases. Another 180 original-target bounded cases passed (148 normal, 32 cluster).
• Cluster repair on the final fork artifact: all 48 CPU-oracle calls passed normally and under each unfiltered sanitizer, with zero memory errors/race hazards/warnings. Normal-path CPU-oracle/padding-canary checks also passed (10 cases / 2 unsupported BF16-sorted skips). No CUDA reference reductions, sanitizer filters or suppressions were used.
• Minimum runtime compatibility: the same final wheel, without rebuilding, passed under Python 3.10.21 / PyTorch 2.10.0+cu130: 24 CPU checks, 52 exact SM120 CUDA output checks and 31 validation checks, including current-stream handling and changed-input graph replay. This is bounded actual execution, not solely an ABI tag or symbol inspection.

Follow-up build cleanup restores upstream setup style and the common --use_fast_math --ftz=false flags, leaving only gencode flags architecture-specific. A fresh SM120/SM121 wheel passed the checker (36 kernel instances, zero STACK/LOCAL), all 15 build-plan tests, the 2,763-pass/198-skip SM120 suite and the same 19-case unfiltered memcheck/racecheck selections. SM120/SM121 normalized SASS and resource tables are identical to the preceding no-fast-math build. The earlier full-mixed and minimum-runtime results above retain their original artifact scope; they are not claimed as fresh runs of this follow-up wheel. README now matches the target dev version byte-for-byte to minimize future upstream-sync conflicts; it is not part of the wheel payload.

SM120 performance versus torch.topk

Complete 54-row results and methodology. Measured on RTX PRO 6000 Blackwell Server Edition (188 SMs), PyTorch 2.13.0+cu130. All 54 pre-timing correctness checks passed. DeepSelect was faster in 15 rows and slower in 39; this does not establish a general speedup.

Speedup below means Torch pipeline latency / DeepSelect latency (>1 favors DeepSelect):

Comparison	Faster rows	Observed paired speedup range
BF16 K512, unsorted values + int64 indices on both sides	4 / 20	0.124–1.217x
BF16 K512, int32 index-only output; Torch topk + cast included	8 / 20	0.391–1.494x
FP32 sorted sampling, width129280/K512	0 / 3	0.325–0.581x

The remaining 11 rows cover generic K1024/4096, the B16→B17 segmented-dispatch boundary, and ascending/tie-heavy distributions. In the boundary case (BF16, width65536, K512, index-only), DeepSelect latency rises from 38.676 µs at B16 to 98.545 µs at B17; the corresponding Torch pipelines take 34.956 and 33.905 µs. Long-row int64 and FP32 sorted paths are also slower and need further optimization.

Timing uses 20 calls per warmed CUDA graph, six alternating paired rounds within one process, and ten CUDA-event samples per round. Reported latencies are medians of round medians; speedups are medians of paired ratios. Allocation during capture, graph construction, CPU validation and input preparation are excluded. Observed ranges are not confidence intervals. Clocks were not forced. This is steady-state GPU pipeline timing, not eager/serving latency or a before/after measurement of the CUB memoization change.

Scope and remaining limitations

• SM121 is explicitly compiled and resource-checked, but no SM121 hardware runtime result is claimed.
• GPU runtime qualification covers the available SM120 and SM103 devices; SM100a compilation is not SM100 hardware qualification.
• The existing vLLM CMake consumer selects only SM100-family sources/architectures. Enabling SM12 in the engine requires a separate consumer change; this PR does not claim to enable it automatically.
• No universal SM12 speedup or zero-regression claim is made. The Torch comparison above measures the current implementation; isolating the generic-sort CUB trade-off still requires a separate before/after comparison.
• Stable ABI and ABI3 do not remove CUDA, driver, glibc or C++ runtime requirements. The tested toolchain/runtime versions and sanitizer scope will be stated explicitly.
• The inherited source-distribution version helper may read an enclosing Git repository's revision when rebuilt inside another checkout. This PR does not change that versioning behavior.

Attribution

The input-device/current-stream and vector-boundary safety work builds on upstream #4 and #6 by @morluto. Existing SM100/SM103 algorithms and the fork's stable ABI work are retained.

[lucifer1004]

SM120 performance comparison with PyTorch — all 54 cases

DeepSelect was faster in 15/54 cases and slower in 39/54. Core BF16 mode A won 4/20, core BF16 mode B 8/20, and the primary FP32 sampling group 0/3. Generic cases contributed the remaining 3 wins out of 6; the boundary case and all 4 supplementary-distribution cases were slower. No cases were discarded. These results do not establish a universal or aggregate speedup.

Important slowdowns versus Torch:

• Mode A, B=1, N=131072, K=512: DeepSelect 234.966 µs versus Torch 29.046 µs, 0.124×.
• Mode B batch-boundary cliff, N=65536, K=512: B=16 is 38.676 vs 34.956 µs (0.904×); B=17 rises to 98.545 vs 33.905 µs (0.344×). Both sides of the boundary are included below.
• FP32 sorted sampling, N=129280, K=512: all three batches are slower, 0.325–0.581×.

Environment and scope

NVIDIA RTX PRO 6000 Blackwell Server Edition, SM120, 188 SMs; Python 3.12.13; PyTorch 2.13.0+cu130, PyTorch CUDA runtime 13.0; installed CUDA toolkit 13.2 (nvcc 13.2.86); NVIDIA driver 595.58.03. Clocks were not forced.

PR source revision: 60f69a1. Exact tested extension SHA-256: 6d5d634782b8b8a7bac677a536bb103aa060fe24b6c88caed748b5feb69350ff.

This is a DeepSelect-versus-Torch comparison, not a before/after test of the CUB memoization change. It measures finite-input, warm, fixed-input steady-state GPU pipeline latency, not allocation/host overhead or serving throughput. It provides no SM121 hardware performance evidence. DeepSelect's always-on NaN checks remain enabled; this is not equal instruction work.

Output contracts and timing

• A: values plus int64 indices, sorted=False on both sides. DeepSelect's int64 path uses the generic implementation, bypassing the specialized int32 fast paths.
• B: int32 indices only, DeepSelect return_value=False versus torch.topk(..., sorted=False).indices.to(torch.int32). Torch timing includes its value computation and int64-to-int32 cast; output requirements match, but implementation work differs.
• Sampling: FP32 values plus int64 indices, sorted=True on both sides. The generic sorted cases use this same contract.

Each CUDA graph contains 20 calls with separate retained outputs. There are 6 newly captured paired rounds in the same process, alternating DeepSelect/Torch then Torch/DeepSelect order. Each implementation has 5 eager warmups and 3 graph warmups per round, followed by 10 CUDA-event samples; each sample is one graph replay's elapsed GPU time divided by 20. Both implementations allocate their outputs during capture and retain those graph-private outputs through replay. Capture, allocation, Python/CPU overhead and correctness checks are excluded from timing. Inputs are fixed across warmups and replays, not a cold-cache workload.

Latency columns are the median of the six per-round medians. Speedup is Torch / DeepSelect, summarized as the median of six paired per-round ratios, with their minimum–maximum in parentheses. Above 1× means DeepSelect is faster; below 1× means slower. The ratio of the two displayed latency medians need not equal the reported median paired ratio. All values are rounded to three decimals; narrow ranges can therefore display equal endpoints. These are same-process observations, not independent process trials; the min–max ranges are not confidence intervals.

54/54 pre-timing correctness checks passed. These checked valid, unique indices, gathered-value bit equality when values were returned, sorted order when requested, and exact equality of selected-value multisets against Torch. Ties need not produce identical indices. This does not mean every timed graph replay was revalidated. Primary inputs are normal-distributed; ascending and tie-heavy inputs are separately identified supplementary cases.

Complete results

B = batch size, N = row width, K = selected elements per row. DS = DeepSelect; Torch = the complete mode-specific Torch pipeline above. Latencies are µs per call. Tables contain 20 + 20 + 3 + 6 + 1 + 4 = 54 cases, each exactly once.

Core A — BF16, normal, unsorted values + int64 indices (20 cases; 4 faster, 16 slower)

B	N	K	DS µs	Torch µs	Speedup × (min–max)
1	8192	512	23.525	25.407	1.080 (1.078–1.082)
1	32768	512	74.880	25.114	0.335 (0.335–0.336)
1	65536	512	128.591	26.814	0.209 (0.208–0.209)
1	131072	512	234.966	29.046	0.124 (0.123–0.124)
4	8192	512	23.635	25.712	1.088 (1.087–1.089)
4	32768	512	75.242	26.738	0.355 (0.355–0.356)
4	65536	512	129.144	28.780	0.223 (0.223–0.223)
4	131072	512	235.802	31.814	0.135 (0.135–0.135)
16	8192	512	23.670	26.220	1.108 (1.106–1.110)
16	32768	512	75.423	28.807	0.382 (0.382–0.382)
16	65536	512	129.311	33.066	0.256 (0.255–0.256)
16	131072	512	236.442	39.041	0.165 (0.165–0.165)
64	8192	512	23.655	28.788	1.217 (1.216–1.218)
64	32768	512	75.559	37.051	0.490 (0.489–0.491)
64	65536	512	129.446	46.969	0.362 (0.362–0.364)
64	131072	512	237.294	66.666	0.281 (0.281–0.281)
256	8192	512	46.482	37.073	0.798 (0.796–0.798)
256	32768	512	149.066	74.594	0.500 (0.500–0.502)
256	65536	512	256.132	121.166	0.473 (0.473–0.473)
256	131072	512	467.617	185.602	0.397 (0.397–0.397)

Core B — BF16, normal, unsorted int32 indices only (20 cases; 8 faster, 12 slower)

B	N	K	DS µs	Torch µs	Speedup × (min–max)
1	8192	512	22.110	26.856	1.215 (1.214–1.216)
1	32768	512	54.688	26.960	0.493 (0.492–0.494)
1	65536	512	37.664	28.759	0.764 (0.763–0.764)
1	131072	512	47.408	30.937	0.653 (0.652–0.653)
4	8192	512	22.099	27.229	1.232 (1.231–1.235)
4	32768	512	54.775	27.264	0.498 (0.498–0.499)
4	65536	512	38.081	29.277	0.769 (0.766–0.771)
4	131072	512	46.502	33.240	0.715 (0.714–0.716)
16	8192	512	22.813	27.802	1.219 (1.217–1.221)
16	32768	512	56.488	30.840	0.546 (0.546–0.546)
16	65536	512	38.676	34.956	0.904 (0.903–0.905)
16	131072	512	47.697	40.768	0.855 (0.854–0.857)
64	8192	512	23.828	30.840	1.294 (1.292–1.296)
64	32768	512	55.973	39.053	0.698 (0.697–0.698)
64	65536	512	99.026	48.668	0.492 (0.491–0.492)
64	131072	512	174.727	68.245	0.391 (0.390–0.391)
256	8192	512	26.121	39.032	1.494 (1.491–1.497)
256	32768	512	61.382	76.217	1.242 (1.241–1.243)
256	65536	512	105.048	122.744	1.169 (1.167–1.170)
256	131072	512	183.925	187.278	1.018 (1.018–1.019)

Primary sampling — FP32, normal, sorted values + int64 indices (3 cases; 0 faster, 3 slower)

B	N	K	DS µs	Torch µs	Speedup × (min–max)
1	129280	512	182.830	59.503	0.325 (0.325–0.326)
16	129280	512	197.270	67.702	0.343 (0.343–0.343)
64	129280	512	199.892	116.208	0.581 (0.581–0.582)

Generic K coverage — normal, B=4, N=8192 (6 cases; 3 faster, 3 slower)

A is unsorted; sampling is sorted. Both return values and int64 indices.

Mode	Dtype	K	DS µs	Torch µs	Speedup × (min–max)
A	BF16	1024	25.025	25.895	1.035 (1.035–1.037)
A	BF16	4096	24.761	26.312	1.063 (1.062–1.066)
A	FP32	1024	23.410	31.574	1.349 (1.348–1.351)
A	FP32	4096	27.721	25.880	0.933 (0.932–0.935)
sampling	FP32	1024	50.476	44.556	0.883 (0.882–0.883)
sampling	FP32	4096	59.202	44.209	0.747 (0.746–0.747)

Batch boundary — mode B, BF16, normal (1 case; slower)

The B=16 counterpart is already included in core B, not duplicated here.

B	N	K	DS µs	Torch µs	Speedup × (min–max)
17	65536	512	98.545	33.905	0.344 (0.344–0.344)

Supplementary distributions — mode B, BF16, N=65536 (4 cases; 0 faster, 4 slower)

Distribution	B	K	DS µs	Torch µs	Speedup × (min–max)
ascending	1	512	35.914	28.762	0.801 (0.800–0.802)
ascending	64	512	104.421	48.841	0.468 (0.467–0.468)
tie-heavy	1	512	35.482	28.772	0.811 (0.811–0.811)
tie-heavy	64	512	95.018	48.551	0.511 (0.510–0.511)