Fix incorrect sparse-attention top-k selection when candidate bins overflow

[ormandj]

Motivation

DSA top-k selection can drop candidates when a coarse radix bin reaches its fixed buffer capacity, selecting lower-scoring tokens before full scores are compared. This affects exact-capacity and oversized bins in KPool JIT, DeepSeek-V4 JIT v1/v2, legacy CUDA AOT and DeepSeek-V4 AOT selectors.

Modifications

Keep normal-bin selection. At the boundary, refine the exact FP32 cutoff and emit higher scores and cutoff ties in a final bounded row scan. Preserve page mapping, -1 padding and synchronization around shared-memory reuse. Retain main's additive ROCm tests. Co-authored with Benjamin Oldenburg (@bold84).

Accuracy Tests

Author validation at 759070d497 on main afe90a8bc9 on one RTX PRO 6000 Blackwell Max-Q (SM120): 297 tests and 38 subtests passed. Targets: test/registered/kernels/ops/attention/test_topk_v2.py, test/registered/kernel/attention/test_dsa_kpool_topk_transform.py. This run exercises source JIT selectors and does not rebuild the CUDA or ROCm AOT binaries.

Author tests at 010c3fd303 on main dcebe8c473 passed 297 JIT tests plus 38 subtests on RTX PRO 6000 Blackwell Max-Q (SM120). This run exercised the source JIT selectors; it did not rebuild the CUDA or ROCm AOT binaries. The additional AOT results below describe an earlier source-built snapshot.

Author SM120 tests of the isolated repaired snapshot passed 12 KPool tests plus 38 subtests, 16 focused JIT tests and 14 source-built CUDA AOT cases. The regressions compare selected score multisets with an exact reference, including ties, signed zero, input permutations, offsets and idle rows. ROCm AOT was not compiled or run locally.

Speed Tests and Profiling

Overflow handling adds row scans; no zero-overhead claim is made. The earlier integrated serving comparison remains in its measurement comment. It was not a kernel-isolated benchmark and was not repeated for this rebase.

Checklist

• Format changed code and retain regression coverage.
• Report author validation and preserve contributor credit.
• Complete upstream CUDA and ROCm CI.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #34282238811
Latest PR Test (Extra): ❌ Run #34282238726
Latest PR Test (AMD ROCm 7.2): ❌ Run #34282238769

[zRzRzRzRzRzRzR]

Issue received, following and test this.

[zRzRzRzRzRzRzR]

This PR seems to have only fixed its copy in the kpool JIT kernel. There are several other copies of the same code affect all dsa model.

[ormandj]

This PR seems to have only fixed its copy in the kpool JIT kernel. There are several other copies of the same code affect all dsa model.

I was working on addressing issues with GLM 5.3 Flash in my SGlang build that lead to crashes directly, but was keeping the breadth smaller to minimize review burden. It sounds like you desire the full and complete fix, so let me evaluate that. Thank you for the feedback.

[ormandj]

@zRzRzRzRzRzRzR Updated the branch in response to the scope issue you identified.

• The kpool JIT, DeepSeek-V4 JIT, legacy AOT, and DeepSeek-V4 AOT copies now use matching package-local shared selectors; the JIT and AOT packages remain independently buildable.
• DeepSeek-V4 v2 register, streaming, and clustered selection now uses the same full-row FP32 fallback when its tie buffer overflows.
• The legacy fallback now includes the exact-capacity boundary after an isolated GPU test failed at exactly 4,096 threshold-bin candidates.
• Cumulative histogram reads are barrier-separated from storage reuse, and all four legacy wrappers synchronize before consuming the selected indices.
• Coverage now exercises every affected entry point, including oversized and exact-capacity bins, negative values, input permutations, higher values already emitted, full-key ties, low-byte separation, offsets, idle rows, the AOT decode path, and raw-index and page-table cluster staging.
• Existing AOT and v2 comparisons now require exact selected-value multisets, while accepting index swaps among equal scores. Stale v2 dispatch cases were replaced with cases for the current 32,768 small-batch cluster floor and 512-row cluster limit; 262,144-token coverage remains at batch 40.

The branch is rebased onto SGLang main at 9ed2721c6d. A downstream integration candidate carrying the selection changes passed the kpool JIT, focused DeepSeek-V4 JIT, and source-built SM120 AOT runs on one Blackwell (sm_120) GPU. CUDA and ROCm CI for this exact commit are pending; ROCm was not compiled or run locally.

[zRzRzRzRzRzRzR]

The fix in this PR is correct, but it costs quite a bit of performance as it stands.

Model output quality before vs. after the fix doesn't seem to differ much on long text (at least when measured end-to-end, e.g. by benchmark scores such as GQPA, need more than 32K context) — an interesting finding.

Maybe there's a better way to fix this? cc @JustinTong0323

[JustinTong0323]

The fix in this PR is correct, but it costs quite a bit of performance as it stands.

Model output quality before vs. after the fix doesn't seem to differ much on long text (at least when measured end-to-end, e.g. by benchmark scores such as GQPA, need more than 32K context) — an interesting finding.

Maybe there's a better way to fix this? cc @JustinTong0323

Kernel level there is some perf regression but not obvious in e2e, so I think it shall be ok.

[JustinTong0323]

@DarkSharpness to review :)

[ormandj]

The fix in this PR is correct, but it costs quite a bit of performance as it stands.
Model output quality before vs. after the fix doesn't seem to differ much on long text (at least when measured end-to-end, e.g. by benchmark scores such as GQPA, need more than 32K context) — an interesting finding.
Maybe there's a better way to fix this? cc @JustinTong0323

Kernel level there is some perf regression but not obvious in e2e, so I think it shall be ok.

I have some ideas on how to optimize this, I'll give it a shot and report back.

[ormandj]

Updated the branch with a cutoff-first oversized-bin fallback.

• The four exact-key refinement passes now build histograms only; one final row scan writes strictly greater values and cutoff ties into disjoint output ranges.
• DeepSeek-V4 reuses the bounded candidates captured by the coarse pass when every value in the oversized bin has the same exact key.
• The legacy AOT/JIT and DeepSeek-V4 selectors use the same cutoff/emission split. Coverage now includes byte-level cutoff separation, exact ties after topk - 1 higher values, signed zero, and all-equal oversized bins.

[ormandj]

Measured PR head 4302a2b719 in our served GLM-5.3-Flash W4A16 profile on 2x RTX PRO 6000 Blackwell (SM120), TP2. C1-C3 used byte-identical fixed-window configs versus the previous served build, with five analyzer-valid repetitions per cell.

Cohort	Previous mean fwd/s	Current mean fwd/s	Mean change	Previous median fwd/s	Current median fwd/s	Median change
C1	62.88	62.29	-0.94%	63.31	63.72	+0.64%
C2	42.92	42.83	-0.21%	43.55	42.66	-2.03%
C3	35.45	35.54	+0.23%	35.57	35.83	+0.76%

The largest mean decrease was 0.94% and the largest median decrease was 2.03%, while the per-cell sample CVs were 3.83% to 6.31%. The current C4 panel also completed five analyzer-valid repetitions at 31.82 mean and 31.98 median fwd/s. This is a serving-level comparison rather than a kernel-isolated A/B, and it does not show a material forward-pass-rate regression with the cutoff-first patch included.

[AranKomat]

We independently reproduced this candidate-buffer overflow on MI300X/gfx942
while qualifying GLM-5.3-Flash's pooled DSA path. Adding AMD evidence here rather
than opening another selector implementation.

In our frozen runtime, exactly 319 of 567 captured score rows both overflowed
the 4096-entry cutoff bucket and failed an independent exact selected-value
oracle. A local bounded full-score rescan repaired all 567 rows in three
repetitions (1701 checks), plus 18 changed-input graph checks on synthetic
crowded-bucket/tie/offset/ragged fixtures. Subsequent live all-rank/all-layer
98K/131K selection and retrieval checks passed.

This supports the failure mechanism and usefulness of the fix on gfx942. It
does not validate this PR's current head: those checks used our independently
patched historical header, not this branch. Equal-score index swaps were
allowed; we did not claim deterministic tie order or exhaustive physical-page
mapping coverage. The repair was common to both arms of our later serving
comparison, so we do not attribute that comparison's speedup to this fix.

The cutoff-first fallback described here looks like the right place to
consolidate this work. Are there particular AMD boundary/graph cases you would
like added to the existing regression suite? We can focus follow-up on testing
the maintainer-preferred implementation rather than maintaining a competing fix.

AI assistance was used to summarize our recorded experiments and prepare this
comment; no new hardware run of this PR is being reported.