[DSA] Fix top-k v2 emitting invalid indices under tie overflow / inf scores (IMA in FA3 sparse decode) - #30645
Conversation
…scores The fused top-k v2 kernel could leave -1 (or short) entries inside the first min(seq_len, topk) output slots whenever the threshold bin held more than kMaxNumTie candidates (e.g. ReLU-degenerate indexer scores) or touched the fp16 +/-inf key space. Downstream FA3 sparse decode reads exactly min(seq_len, topk) slots, so a -1 there dereferences kv_pool_base - row_bytes and crashes with an illegal memory access (nightly GLM-5.2-FP8 perf test, coredump faulting address = base - 1152). - Raise kMaxNumTie to 2048 (= kMaxTopK): the tie stage may need to fill up to topk slots (above_count can be 0), so any smaller cap leaves a hole. Overlay the tie machinery with the dead coarse histogram so the bigger buffer shrinks the shared-memory footprint (24.9KB -> 18.9KB). - Split the exact radix tie-select into a <1>-item common path (num_ties <= kBlockSize, identical to the old code) and a <2>-item overflow path, and fix the tie emit / cluster tie merge loops that silently assumed num_ties <= kBlockSize. - Make coarse_bin_lower_bound inf-aware: treat the +/-inf ordered keys as +/-65536 (midpoint lands on the exact fp32->fp16 RN overflow threshold +/-65520) and saturate NaN-space keys, so thresholds at or next to the inf bin classify +/-inf and >65504 scores instead of collecting nothing. Hoist the finite fast path above the special cases (verified bit-exact over every bin for kBits 10 and 12). - Pad unfillable slots with a valid in-range index instead of -1 as a last-ditch backstop (unreachable for NaN-free scores): duplicates cost accuracy, never an IMA. Verified on H200: selected-score sets match torch.topk exactly across random / fat-tie / all-equal / inf-heavy / >65504 / -inf cases on all three kernel paths (register, streaming, cluster); compute-sanitizer memcheck and racecheck clean; marker benchmark shows healthy paths within noise and the degenerate fat-tie path +5-8% (it now exact-ranks 2048 tie candidates instead of dropping half). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-run-ci-label |
|
/rerun-test test_glm52_fp8.py |
|
Results for 🚀 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes and refactors the TopK kernel implementation for DeepSeek v4. Key changes include adding a fast path to coarse_bin_lower_bound, increasing kMaxNumTie to 2048 to avoid downstream illegal memory accesses, refactoring handle_tie and radix_tie_select to support multi-item register layouts, and overlaying shared memory structures to control footprint size. The review feedback highlights two important issues: first, a potential uninitialized memory read in radix_tie_select when topk is 0, which can be resolved with an early return in handle_tie; second, a strict aliasing violation in to_finite_val due to pointer-based reinterpretation, which should be replaced with __ushort_as_half.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| @@ -208,15 +244,11 @@ struct TopKConfig { | |||
| static_assert(kNumWarps == kWarpSize); | |||
|
|
|||
| if (num_ties <= topk) { | |||
There was a problem hiding this comment.
If topk is 0 (which happens when above_count >= global_topk), handle_tie is still called. If num_ties > kWarpSize * 4, it will enter radix_tie_select with topk_remain = 0. In radix_tie_select, since topk_remain is 0, the condition above < topk_remain is never met, meaning no thread writes to smem->match. This causes threads to read uninitialized/garbage values from smem->match and perform atomic operations/scatters based on garbage threshold bins. Adding an early return if (topk == 0) return; at the beginning of handle_tie completely avoids this issue and saves unnecessary work.
if (topk == 0) return;
if (num_ties <= topk) {
| const auto to_val = [](uint32_t okey) -> float { | ||
| // ordered16 -> fp16 value (inverse of the transform in extract_coarse_bin); | ||
| // finite keys only. | ||
| const auto to_finite_val = [](uint32_t okey) -> float { |
There was a problem hiding this comment.
Using reinterpret_cast to cast a uint16_t* to const fp16_t* and dereferencing it violates strict aliasing rules in C++ and can lead to undefined behavior. Since hb is a uint16_t and fp16_t is __half, you can use the standard CUDA/HIP intrinsic __ushort_as_half(hb) to safely perform the bit-reinterpret cast without pointer dereferencing.
|
/rerun-test test_glm52_fp8.py |
|
Results for 🚀 🚀 |
|
/rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_hisparse.py test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py test/registered/kernels/test_dsa_indexer.py test/registered/jit/deepseek_v4/test_topk_v2.py |
|
Results for 🚀 🚀 🚀 |
|
@DarkSharpness @Fridge003 @b8zhong This PR broke the lint... |
…indices under tie overflow / inf scores (IMA in FA3 sparse decode) (sgl-project#30645) (sgl-project#30698) Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…sgl-project#30699) (sgl-project#30704) Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
…scores (IMA in FA3 sparse decode) (sgl-project#30645) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
4 participants
(generated by claude)
Motivation
The nightly
test_glm52_fp8.pyperf test crashed with an illegal memory access on all three variants (run). The CUDA coredump shows the faulting kernel is FA3 sparse decode (FlashAttnFwdSm90, bs=1 decode right after an 8192-token prefill), and the faultingLDGSTSaddress is exactlykv_pool_base - 1152B(= one 576-elem bf16 DSA KV row) — i.e. FA3 dereferenced KV slot index-1.The
-1comes from the fused top-k v2 kernel. FA3 reads exactlymin(seq_len, topk)entries of the top-k output (dsa_cache_seqlens = seq_lens.clamp(max=topk)), so every one of those slots must hold a valid index. Two data-dependent extremes break that contract:kMaxNumTie = 1024candidates from the threshold bin. When the bin holds more (ReLU-degenerate indexer scores easily put thousands of positions on one exact value — the random-token perf workload hits this; short gsm8k requests never reach the sparse path, which is why accuracy passed) andtopk - above_count > 1024,handle_tiepads the remaining slots with-1. Unit repro on H200: seq=8320, topk=2048, 100 winners + 8220 zeros → 924 invalid entries inside the valid prefix.coarse_bin_lower_boundcomputed NaN boundaries for bins touching the fp16 ±inf keys (their neighbors live in NaN key space), so the collect pass matched nothing; large-but-finite scores (>65504, which fp16-round into the inf bin) were dropped entirely.Modifications
All in
topk_impl.cuh(used by the v2 JIT kernel only; v1 is self-contained):kMaxNumTie1024 → 2048 (=kMaxTopK), with astatic_assertpinning the invariant. 1536-style middle values don't close the hole:remain = topk - abovereachestopkitself wheneverabove == 0(all-equal / inf-heavy rows). The tie buffer is overlaid with the coarse histogram (dead oncefind_thresholdpublishes the bin;TieHandleSmemandtie_valuesstay side by side since they are live together), so per-block shared memory shrinks from ~24.9KB to ~18.9KB.radix_tie_select<1>for the commonnum_ties <= kBlockSizecase (instruction-identical to the old code),radix_tie_select<2>only for the rare overflow case. Also fixes the tie emit and cluster cross-rank tie merge loops, which silently assumednum_ties <= kBlockSize.coarse_bin_lower_bound: ±inf ordered keys act as ±65536 so the midpoint lands exactly on ±65520 (the fp32→fp16 RN overflow threshold); NaN-space keys saturate. Boundaries stay monotone and never NaN. The finite fast path is hoisted above the special cases and range-checks both keys at once — verified bit-exact against the general path for every bin of kBits 10 and 12.-1— a duplicate token costs a little accuracy, never an IMA.Verification (H200, same hardware as the failing job)
torch.topkexactly across random / fat-tie / all-equal / inf-heavy / inf-some / huge-finite(1e30) / -inf-heavy distributions, on all three kernel paths (register, streaming, cluster); no duplicate or out-of-range indices; the original IMA repro (including all-NaN) now yields only in-range indices.compute-sanitizermemcheck and racecheck clean (racecheck also validates the smem overlay phase separation).do_bench, CUDA-graph replay, cold L2) A/B vs main: healthy-distribution shapes within ±1% (cluster small-batch +~2%); the degenerate fat-tie path +5–8%, which is the cost of exact-ranking 2048 tie candidates instead of dropping half of them.Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #29012052859
Latest PR Test (Extra): ❌ Run #29012052641