feat(topk): add sample-verified exact Top-K

[VAthree]

Summary

This PR adds hpc.topk_filtered, an exact FP32 row-wise Top-K operator for long-context sparse-attention indexers on NVIDIA Hopper GPUs.

Top-K is a key operator in the sparse-attention indexers used by models such as GLM-5.3 and Hy4-preview. An indexer first produces a score row over the available KV positions, selects a fixed number of indices, and passes those indices to sparse attention. As context length grows, the downstream attention work remains bounded by K, while standalone Top-K must still process every valid score in the row.

HPC-Ops Top-K uses a small, regular view of the current row to propose a coarse upper-tail boundary before the mandatory complete-row traversal. The complete traversal validates that proposal, forms the candidate set, and starts exact FP32 refinement in one fused pass. Current-row sampling thereby removes the initial histogram-then-classification reread from the common long-row path, while complete-row validation preserves exact output semantics.

Design

A conventional radix selector first histograms the complete row to discover its rank boundary. Applying that boundary then requires a second traversal or a row-sized retained copy.

Early filtering has a coarser objective than final selection: retaining a compact upper tail containing the answer. A regular 1/s view of the current row provides this signal, with a deterministic row-dependent phase distributing consecutive rows across the s possible offsets. Its order statistic locates a conservative O(K) upper-tail boundary. The sample controls the retained workload; the complete row certifies candidate sufficiency and the FP32 refinement determines the returned indices. On the common sampled path, score-row traffic is approximately one fractional view plus one complete traversal before frontier-only refinement.

Dual-threshold rescue. The coarser wide-row view records two nested boundaries during the same sampled histogram prefix. The primary boundary serves the common path; a deeper secondary boundary remains dormant as a bounded recovery margin. Both boundaries share the same view and prefix scan.

The operator follows three logical phases:

1. Coarse boundary localization. Sampled FP32 scores are projected into a monotone 11-bit coarse ordering and accumulated into an on-chip histogram. An upper-tail order statistic proposes a row-local primary boundary with a compact retention margin. The wide-row policy also records the nested secondary boundary from the same histogram.
2. Fused validation and candidate formation. The proposed coarse boundary is converted once to an inclusive FP32 cutoff. One complete-row traversal validates the admitted count, persists admitted indices, and constructs the first exact FP32 radix histogram on the candidates' first encounter. If the primary tail is underfilled, a recovery traversal appends the band newly admitted by the secondary boundary and extends the existing histogram. Continued underfill enters sample-independent exact coarse recovery.
3. Exact FP32 refinement. The retained upper tail is refined with an 11+11+10 decomposition of the complete FP32 ordered key. Each round commits groups strictly ahead of the rank boundary, discards groups behind it, and carries the unique boundary group to the next digit.

Sampling proposes the workload-reduction boundary; complete-row validation certifies candidate sufficiency, and every continuation finishes with the original FP32 values. An underfilled primary proposal may use the nested recovery boundary first, while an unavailable view or continued underfill uses the sample-independent exact coarse route. The operator returns an unordered Top-K set with an arbitrary valid choice among equal-valued boundary indices.

GPU implementation

The primary long-row path uses a persistent pool of 512-thread CTAs. Each CTA owns the three phases for one row and obtains subsequent rows from a device-side completion-order queue. This balances causal and ragged row lengths while keeping launch dimensions stable for CUDA Graph replay.

The sampled path combines vectorized cache-global loads, FP32 cutoff classification, candidate persistence, and construction of the leading exact histogram. Exact refinement revisits the unresolved frontier through ping-pong candidate buffers, and the shared 2,048-bin histogram allocation is reused across phases and radix digits. Persistent workers own bounded overflow slices, bounding scratch usage by the worker pool.

For a small number of long rows, a KV-split mapping partitions the complete-row validation across cooperating CTAs and merges compact count and histogram state before one finisher performs refinement. Shorter rows and a few other complementary shapes use direct-exact row-local or cooperating-CTA mappings. All mappings share the same operator contract and FP32 refinement procedure.

API

counter_bytes, workspace_bytes = hpc.topk_filtered_workspace_size(
    num_rows=logits.shape[0],
    max_kv_len=logits.shape[1],
)

counters = torch.zeros(counter_bytes, dtype=torch.uint8, device="cuda")
workspace = torch.empty(workspace_bytes, dtype=torch.uint8, device="cuda")

hpc.topk_filtered(
    logits,           # [M_cap, N], float32
    ke,               # [M_cap], int32 valid-prefix lengths
    output,           # [M_cap, K], int32
    num_valid_rows,   # one-element device int32 scalar M_live
    top_k,
    counters,
    workspace,
)

M_cap is the captured row capacity and M_live <= M_cap is the runtime number of valid rows. Row r selects from [0, ke[r]); physically padded columns may contain arbitrary values.

The counter buffer must be zero-initialized before first use and is returned to the same state after every launch. The candidate workspace needs no initialization and can be retained across CUDA Graph replays.

Performance

We randomly sample chunks with distinct shapes to cover a variety of real chunking patterns. K is fixed at 2,048, while the valid KV length varies by row; M denotes the number of rows and N denotes the allocated row length.

The values below are median device latency in milliseconds over 50 iterations on a single NVIDIA H20 with 78 SMs, measured from Nsight Systems NVTX projections. All compared implementations perform exact Top-K selection: no unselected score exceeds any selected score.

Implementation	M=5,536 N=128K	M=1,536 N=192K	M=5,536 N=256K	M=3,136 N=320K	M=1,536 N=384K	M=43 N=420K
HPC-Ops Top-K	1.305	0.531	2.344	1.631	0.963	0.062
SGLang JIT Top-K v2	1.839	0.804	3.527	2.508	1.509	0.073
TensorRT-LLM decode_varlen	2.161	0.921	4.261	3.017	1.825	0.117
TensorRT-LLM prefill GMEM_SPILL	2.112	0.869	3.854	2.862	1.705	0.200
vLLM top_k_per_row_prefill	2.297	0.988	4.419	3.151	1.895	0.253
TensorRT-LLM prefill REREAD	2.174	0.919	4.271	3.035	1.839	0.411
vLLM large_context_topk	3.170	1.353	6.092	4.368	2.590	0.262
vLLM top_k_per_row_decode	3.254	1.324	13.710	8.468	4.573	0.161
TensorRT-LLM single_pass_multi_cta	13.204	4.539	16.197	11.934	5.935	0.185

The SGLang row uses its raw-index JIT Top-K v2 entry, topk_transform_512_v2, evaluated at K=2048 with page-table transformation disabled.

HPC-Ops Top-K has the lowest latency for all six shapes and is 1.17x to 1.57x faster than SGLang, the closest competitor. These measurements report internal Top-K latency.