Fix clipped stage-1 bins and bound kpool top-k selection stores

[ormandj]

Base commit: cf0f761.

Motivation

kpool_topk_transform_kernel<K> (python/sglang/kernels/jit/csrc/dsa/kpool_topk_transform.cuh) selects the top K pool groups of a score row with a two-stage radix select and expands them to token indices for the DSA kpool indexer.

Reported production failure: on a GLM-5.3-Flash W4A16 TP2 deployment on two RTX PRO 6000 Blackwell GPUs, a long-prefill workload produced Xid 31 on both ranks in the same second; the driver coredump named kpool_topk_transform_kernel<512> with a shared-memory store at offset 39,204 against a 39,184-byte window, i.e. stash position 4101 of the second stash buffer. Reported by @bold84 in ormandj/sglang-glm53-flash-sm120 issue #3 and PR #4, and reproduced on the reporter's deployment and on a second W4A16 TP2 RTX PRO 6000 deployment. How the base kernel reaches that position from a valid input is not established: under the source-level invariant every threshold finder fires, so every stash counter is reset each round and every stash position is bounded by the stash count. A compute-sanitizer reproducer on the affected architecture is still to be provided; this PR does not claim to have reproduced the fault.

Selection defect verified from the base source: rows whose stage-1 threshold bin holds more than 4096 candidates (tied or tightly clustered scores) select the wrong groups. Author report from the investigation: on the diverging input (9407 candidates in one bin), 286 of the 512 selected groups differed from torch.topk.

Modifications

• Stage-1 bins larger than the stash were clipped by atomic arrival order: the threshold bin's candidates are stashed in a 4096-entry shared buffer, atomicAdd kept the first 4096 by arrival and the guarded store dropped the rest, so the refine rounds ranked a subset. Now, if the stage-1 threshold bin holds more than the stash, the kernel fills the definite members above the bin, then walks key bytes 16, 8 and 0 restricted to the chosen byte path, recomputing histogram, cumsum and threshold at each level and filling that level's above-threshold members, until the remaining bin fits the stash or the last byte is reached, where the clip drops exact ties only.
• All stages bin on bytes of the FP32 monotone key; the unused FP16 conversion and its include are removed. The base's mixed FP16 coarse key and FP32 refine keys were not a defect (FP16 rounding is monotone); the re-keying lets the overflow descent walk one key, at the cost stated below.
• State initialization is unconditional at each point where the value is consumed: stage 1 resets the threshold-bin id, the first stash count and the output count; each descent finder resets the threshold-bin id, and a descent level resets the stash count before filling the selected bin; each refine round resets the next-round stash count, the final-round remaining count and the threshold-bin id. A finder with no satisfier leaves the threshold at -1, for which the guarded bin-population calculation returns 0.
• Bounded stores: every index[pos] store is guarded by pos < K (the round-3 fill by pos > 0 && pos <= K), stash stores by pos < SMEM_INPUT_SIZE; static_assert(K < SMEM_INPUT_SIZE).
• Sentinel slots: index[] is initialized to -1 before selection, and the expansion writes -1 and skips the page-table lookup for group_id < 0.
• Score contract: scores within the row must be finite; NaN ordering is unspecified (the key orders a NaN by its sign bit, as the base kernel and deepseek_v4/topk_v1.cuh do). The indexer produces FP32 relu-weighted FP8 dot products with -inf only as padding outside the row.

Author report: the fix was contributed by @bold84 in ormandj/sglang-glm53-flash-sm120#4.

Accuracy Tests

test/registered/kernels/test_dsa_kpool_topk_transform.py (register_cuda_ci, base-b-kernel-unit, 1-gpu-large, pool_size = 4, group top-k 512 and 256). Each row is compared with torch.topk by sorted value multiset (any subset of groups tied at the K-th value is accepted); expanded columns are checked to be group_id * pool_size + slot, distinct, below length, and never the -1 pad; rows are padded with +inf beyond lengths. Cases: clusters of 9407, 17802 and 4097 entries in one stage-1 bin, which exercise the clipped-selection path; a 17802-entry row batched with a K + 1 row; 4097 and 8193 identical values; 4097/4097 clusters first differing at key bit 16, 8 or 0; exactly K distinct values above a 65,536-way tie; length == K + 1; uniform rows of several lengths; expansion through a biased page table and through topk_indices_offset. The test proves the clipped-selection defect; it does not exercise the reported out-of-bounds store.

Author report, one RTX PRO 6000 Blackwell (SM120, fresh JIT cache): on this branch the file passes (8 passed, 28 subtests); with the base kernel (cf0f761) substituted, 10 of 16 cases fail on the clustered rows (wrong groups selected), the rest pass. compute-sanitizer --tool memcheck on test_stage1_bin_exceeds_stash reports 0 errors for both the base and this branch, so the reported out-of-bounds store is not reproduced by these inputs on this device. The following checks were also run: clang-format 20.1.7 with the JIT .clang-format, black, isort, ruff, codespell, scripts/lint/check_registered_tests.py; the constructed rows were checked on CPU to land in a single stage-1 bin; the header compiles as a CUDA translation unit with the JIT's flags for SM120 for both K values (nvcc -c, no execution).

Contributor-reported on ormandj/sglang-glm53-flash-sm120#4, not re-run here: exact value-set agreement with torch.topk over 12,000 randomized structured rows; compute-sanitizer --tool memcheck with 0 errors on tied, clustered, boundary and page-table cases including the issue-#3 reproducer; the previously faulting workload (four 150k to 170k-token prefills followed by a 173,504-token restore with concurrent decode, TP2) passed 3/3.

Speed Tests and Profiling

The radix path executes six additional block-wide barriers per row. Rows whose stage-1 bin exceeds the stash run a histogram pass, a fill pass and a stash pass per descent level instead of clipping the bin. Using FP32 key byte 24 for stage 1 creates coarser initial bins than the base's FP16 top byte: for uniformly distributed [0, 1) scores about half the values share the [0.5, 1) bin and exceed the 4096-entry stash at row lengths above about 8192.

Author report, microbenchmark on one RTX PRO 6000 Blackwell (8 rows per call, pool_size 4, topk 2048, fp32 scores, 50 timed calls after warmup), base cf0f761 against this branch:

length	scores	base	this branch
8192	uniform	11.2 us	12.2 us
8192	normal	11.0 us	11.3 us
32768	uniform	22.1 us	28.3 us
32768	normal	18.7 us	18.3 us
131072	uniform	55.3 us	82.0 us
131072	normal	44.8 us	49.5 us

Uniform scores take the overflow descent from 8192 entries on; normally distributed scores mostly do not. Real indexer logits were not benchmarked.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Prepared with AI assistance.

[ormandj]

#37477 merged and closed this PR with it. Re-filed against main as #37625 (same commit, cherry-picks without conflict).