[DSA] Make the fused top-k exact on an overflowing threshold bin - #37941
xiaofei-zheng wants to merge 2 commits into
Conversation
8a746f3 to
3dbc304
Compare
xiaofei-zheng
changed the title
ae5db33 to
7a83d7f
Compare
|
@HaiShaw could you add the The gate currently exits with For context: this is a correctness fix in the fused DSA indexer top-k, found while profiling our own GLM-5.2 decode runs on MI355X. When a threshold coarse bin holds more than |
|
@BBuf @yuan-luo @DarkSharpness please review |
|
/tag-and-rerun-ci |
DarkSharpness
left a comment
DarkSharpness
left a comment
There was a problem hiding this comment.
Some comments on performance
| // Early-out for a bit-identical candidate set: the radix passes cannot | ||
| // separate it and would only re-derive what phase 3 already staged. One | ||
| // distinct exact key means the candidates are interchangeable, so any | ||
| // kMaxNumTie-subset is correct, including the one already in tie.values. | ||
| // Comparing min against max of the key is exact for any value, and costs | ||
| // the one scan that non-degenerate overflow rows pay on top. | ||
| { | ||
| uint32_t key_min = 0xFFFFFFFFu; | ||
| uint32_t key_max = 0u; | ||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| key_min = min(key_min, key); | ||
| key_max = max(key_max, key); | ||
| } | ||
| }); | ||
| // Reduce per-thread extrema: two atomics per thread, not per candidate. | ||
| if (tx == 0) { | ||
| handle->histogram[0][0] = 0xFFFFFFFFu; | ||
| handle->histogram[0][1] = 0u; | ||
| } | ||
| __syncthreads(); | ||
| atomicMin(&handle->histogram[0][0], key_min); | ||
| atomicMax(&handle->histogram[0][1], key_max); | ||
| __syncthreads(); | ||
| const bool bit_identical = handle->histogram[0][0] == handle->histogram[0][1]; | ||
| __syncthreads(); // all threads read the scratch before the loop clears it | ||
| if (bit_identical) { | ||
| // equal_count > kMaxNumTie on entry, so phase 3 filled the whole buffer. | ||
| const auto above_count = smem->count_gt; | ||
| const auto remain_topk = above_count < topk ? topk - above_count : 0; | ||
| handle_tie(smem->tie.values, problem, above_count, kMaxNumTie, remain_topk, handle); | ||
| return; | ||
| } | ||
| } |
There was a problem hiding this comment.
Why do we need this pass? It's not common that all keys are equal in fp32. We should optimize for hot path. This actually adds to additional linear scan cost in normal pass.
| for (uint32_t round = 0; round < 4 && cand_count > kMaxNumTie && remain > 0; ++round) { | ||
| const uint32_t shift = 24 - round * 8; | ||
|
|
||
| if (tx < kRadixSize) handle->histogram[0][tx] = 0; | ||
| __syncthreads(); | ||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| if ((key & mask) == prefix) atomicAdd(&handle->histogram[0][(key >> shift) & 0xFFu], 1); | ||
| } | ||
| }); | ||
| __syncthreads(); | ||
|
|
||
| refine_find_threshold(cand_count, remain, handle); | ||
| const auto match = handle->match; | ||
|
|
||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| if ((key & mask) == prefix && ((key >> shift) & 0xFFu) > match.bin) { | ||
| const auto pos = atomicAdd(&smem->count_gt, 1); | ||
| if (pos < topk) [[likely]] | ||
| problem.emit(pos, idx); | ||
| } | ||
| } | ||
| }); | ||
|
|
||
| prefix |= match.bin << shift; | ||
| mask |= 0xFFu << shift; | ||
| remain -= match.above_count; | ||
| cand_count = match.equal_count; | ||
| __syncthreads(); // `match` is read above and overwritten by the next pass | ||
| } |
There was a problem hiding this comment.
Is this efficient enough? This would involve 8 linear pass. Please reduce that to at most 4 passes.
| if (tx == 0) smem->count_eq = 0; | ||
| __syncthreads(); | ||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| if ((key & mask) == prefix) { | ||
| const auto slot = atomicAdd(&smem->count_eq, 1); | ||
| if (slot < kMaxNumTie) smem->tie.values[slot] = {val, idx}; | ||
| } | ||
| } | ||
| }); |
There was a problem hiding this comment.
same, this passes should be eliminated i guess
DarkSharpness
left a comment
DarkSharpness
left a comment
There was a problem hiding this comment.
Some comments on performance
7e0bc53 to
c05fe16
Compare
|
@DarkSharpness thanks — all three addressed in c05fe16, and measured on MI355X. The min/max early-out is gone. You are right that it was paying a scan on every refinement to serve a degenerate case. Without it the radix rounds still terminate correctly on an all-equal row: they consume all 32 bits and then truncate a bit-identical set, which is exact. Removing it also freed the tie staging buffer, which this path does not need — it re-derives the candidates rather than trusting what the collect pass staged there — so the buffer now carries the histogram. Down to at most four passes. Two changes got there. The radix is 12 bits instead of 8, using that freed buffer (it is exactly The staging pass is eliminated too, fused into the last round's emit the same way. Kernel time, 200 iterations after warmup:
The hot path is unchanged, as it should be — it still only pays the On correctness after the rewrite: the suite is 288/288 on this branch. Run against the unfixed kernel on main it is 8 failed / 2 passed, the two passes being the |
c05fe16 to
fcd3905
Compare
|
@DarkSharpness when you have a moment — the performance changes you asked for are in The branch has also been rebased onto latest
The remaining red checks are unrelated to this diff, as far as I can tell:
Happy to dig into any of those if you would rather they were green first. |
The coarse bin comes from the top bits of the fp16 cast, so one bin spans a quarter-binade. Phase 3 stages threshold-bin candidates under `if (count_eq < kMaxNumTie)` and phase 4 clips `min(equal_count, kMaxNumTie)`, so an overflowing bin leaves `handle_tie` ranking an arrival-order subset and the kernel emits a wrong selected-value multiset with no diagnostic. Found while profiling GLM-5.2 decode on MI355X. The same defect is reported independently from the CUDA side in sgl-project#35257, which quotes the same two sites; referenced here rather than closed by this commit. Refine the overflowing bin on the exact key instead: up to four 8-bit radix passes over the order-preserving key from `extract_exact_bin`, exiting once the refined set fits the staging buffer or once all 32 bits are consumed, where the survivors are bit-identical and truncating them is legitimate. A min/max reduction detects a bit-identical set up front and skips the refinement. The fast path branches on `equal_count > kMaxNumTie` and is otherwise unchanged; kMaxNumTie, kBlockSize and the shared-memory footprint are untouched. Also pin the ROCm `__launch_bounds__` occupancy argument to the physical wave floor: HIP reads it as waves per SIMD where CUDA reads it as blocks per SM, so passing kOccupancy asks for 2 waves/SIMD, below the 4 that a 1024-thread block already forces. Codegen-neutral on its own, but the wrong reading (8 waves/SIMD) caps the allocator at 64 VGPRs, which binds nothing at main's 57 and spills at the 77 this change needs. The pre-existing tests draw from `torch.randn` and carry MAX_PERMIT_ERROR = 5, so they pass unchanged on the truncating kernel. Added narrow, narrow_bin and tiny distributions that collapse a row into one coarse bin, on a register-path and a streaming-path shape, asserting with no tie tolerance. On the unfixed kernel the 6 new cases fail with 227-9198 wrong selections each while the 2 all_equal controls pass; all 286 cases pass with this change. A separate 50-case gate moves from 29/50 to 50/50. Performance is parity: a paired decode A/B measures -0.50% / -0.27% step time against a 0.04-0.08% within-arm spread. On gfx950 the kernels go from 57 to 77 VGPRs with zero spills. Known gap: the CUDA-only `TopKCluster` path still truncates. Refining there needs cluster-wide rather than block-local histogram and emit counters, so it is left for a follow-up. Co-authored-by: Cursor <cursoragent@cursor.com>
Review feedback: the refinement cost up to ten full passes over the input -- one min/max scan, four rounds of histogram + emit, one staging scan. Drop the bit-identical early-out. It was a pure optimization for a degenerate case; without it the radix rounds still terminate correctly on such a row, consuming all 32 bits and then truncating a bit-identical set, which is exact. Removing it also frees the tie staging buffer -- this path re-derives the candidates rather than trusting what the collect pass staged there -- to carry the histogram. Widen the radix from 8 to 12 bits by putting the histogram in that buffer. It is exactly kRefineSize uint32 wide, so the shared-memory footprint is unchanged, and 12 bits cover the 32-bit key in three rounds instead of four. Fuse each round's emit pass with the next round's histogram pass, and the last round's emit with the survivor staging pass: one scan classifies a candidate as above / equal / below the round's threshold digit and does the emit and the histogram (or the staging) together. Worst case is now four passes -- one to seed the first histogram, then one per round. A row whose bin resolves in the first round costs two. Measured on MI355X: the hot path is unchanged (5.596 vs 5.604 us at 6x8192, 12.475 vs 12.507 at 4x32768) and the refinement path drops 29-35% (narrow_bin 34.825 -> 22.482 and 75.559 -> 53.223; two_values 53.509 -> 36.354 and 168.130 -> 119.571). Add a two_values distribution to the exactness tests: two adjacent fp32 values at roughly half the row each, so the refinement has to consume all 32 bits and still has more than kMaxNumTie survivors at that point. That is the key-exhausted exit, now with nothing in front of it, reached from a row that is not degenerate -- all_equal reaches the same exit but any subset of it is trivially correct. Co-authored-by: Cursor <cursoragent@cursor.com>
fcd3905 to
9a57b69
Compare
|
@HaiShaw thanks for the earlier approval. Could you help nudge this one along, or advise on how to proceed? @DarkSharpness requested changes on Sep 9 with three performance comments. All three are addressed in Every CI job that exercises this kernel is now green on both vendors, and |
|
Hi. Could you please temporarily hold on until #38798 is merged? Sorry we're actively working on that and it should land within next few days. I will check this PR in detail later. This topk kernel is very sensitive in register usage on CUDA, and any seemingly irrelevant test might lead to wild performance data, so I will take a very detailed look at this PR later. Please stay tuned. Thanks :) |
…verflowing threshold bin The coarse histogram bin is derived from the top bits of the fp16 cast of each score, so one bin spans about a quarter-binade. Phase 3 stages threshold-bin candidates into the fixed kMaxNumTie = 2048 staging buffer and phase 4 clips tie_count = min(count_eq, kMaxNumTie), so once a threshold bin holds more than 2048 elements handle_tie ranks whichever arrived first and the kernel returns a wrong selected-value multiset with no diagnostic. This is the shared DSA indexer top-k used by DeepSeek-V4.1 decode, prefill and the DeepGEMM candidate path, and it fails silently: every emitted index is a valid KV position, so nothing downstream detects the wrong selection. Refine the overflowing bin on the exact key instead of truncating it: up to three radix rounds (12/12/8 bits) over the order-preserving key from extract_exact_bin. It exits once the refined set fits the staging buffer, or once all 32 bits are consumed -- the key is injective on fp32 bit patterns, so the survivors are bit-identical by then and truncating them is legitimate. The fast path branches on count_eq > kMaxNumTie and is otherwise unchanged; kMaxNumTie, kBlockSize and the shared-memory footprint are untouched (the refinement histogram overlays the dead tie staging buffer in a union). This is a correctness fix, not a speedup: the hot path is unchanged. The CUDA-only TopKCluster path still truncates; refining there needs cluster-wide rather than block-local histogram and emit counters, and is left for a follow-up. Ported from sgl-project#37941, which does not apply cleanly after sgl-project#38829/sgl-project#39098 refactored this file. Identifier adaptations: smem->tie.handle -> smem->tie_handle, smem->tie.values -> smem->tie_values, smem->tie.refine_hist -> smem->refine_hist; the local sites name the counters count_gt/count_eq rather than above_count/equal_count. The PR's ROCm __launch_bounds__ hunk is intentionally dropped (CUDA-only deployment). Verified on H20 (SM90, cc 9.0): - The new test fails 8/10 on the unpatched kernel (the two all_equal controls pass, pinning the fault to truncation) and passes 10/10 patched. - test/registered/kernels/ops/attention/test_topk_v2.py: 298 passed. - cuobjdump resource usage and sizeof(Smem) are byte-identical to baseline: 31-32 registers, zero spills, shared memory unchanged.
|
Sounds good, we'll hold. Two things worth flagging while we wait. The truncation this PR fixes is unchanged in #38798 — the And please don't try to carry this change into #38798; it is safer for us to re-port it afterwards. The refinement re-reads candidates from global memory and relies on On your register point: we only have MI355X here, so every number in this PR is gfx950 (57 to 78 VGPRs, no spills, LDS unchanged). We cannot measure the CUDA side at all. If that is the main concern, it would help if someone with an NVIDIA box could check it — or tell us which shapes and metric you want and we will get them into the PR another way. Ping me when #38798 lands and I'll rebase, re-measure and re-request review. |
4 participants
Motivation
The coarse histogram bin comes from the top bits of the fp16 cast of each score, so one bin spans a quarter-binade. Phase 3 stages threshold-bin candidates under
if (count_eq < kMaxNumTie)and phase 4 clipstie_count = min(equal_count, kMaxNumTie), so once a threshold bin holds more than 2048 elementshandle_tieranks whichever arrived first and the kernel returns a wrong selected-value multiset with no diagnostic. This is the shared DSA indexer top-k, serving every model routed throughsrt/layers/attention/dsa.Found while profiling GLM-5.2 decode on MI355X: device-side counters over an agentic serving run enter the truncation path on ~0.04% of rows across all 8 ranks (556–598 of 1.1–1.5M per rank), largest threshold bin 4139 against the 2048 cap, none of those rows bit-identical. #35257 reports the same defect independently from the CUDA side — referenced, not closed.
Modifications
Refine the overflowing bin on the exact key instead of truncating it: up to four 8-bit radix passes over the order-preserving key from
extract_exact_bin. It exits once the refined set fits the staging buffer, or once all 32 bits are consumed — the key is injective on fp32 bit patterns, so the survivors are bit-identical by then and truncating them is legitimate. A min/max reduction skips the refinement when the set is already bit-identical.The fast path branches on
equal_count > kMaxNumTieand is otherwise unchanged;kMaxNumTie,kBlockSizeand the shared-memory footprint are untouched.Folded in from #37942 (closed): the ROCm
__launch_bounds__occupancy argument now states the physical wave floor, since HIP reads it as waves per SIMD where CUDA reads it as blocks per SM. Codegen-neutral alone; it matters because the plausible wrong reading caps the allocator at 64 VGPRs and would spill at the 77 this change needs.Accuracy Tests
Pre-existing cases draw from
torch.randnand tolerateMAX_PERMIT_ERROR = 5, so they never reach the overflow and pass on the truncating kernel. This adds narrow, narrow_bin and tiny distributions that collapse a row into one coarse bin, on a register-path (6×8192) and a streaming-path (4×32768) shape, asserting exactly. Against the unfixed kernel:all_equalcontrolsall_equaloverflows too, but bit-identically, so any subset is correct and it must keep passing — that asymmetry pins the fault to truncation rather than mis-binning. All 286 pass with this change; a separate 50-case gate moves 29/50 to 50/50.Benchmarking and Profiling
MI355X (gfx950), weighted J over the captured GLM-5.2 decode shapes, graph-replayed:
A paired decode A/B measures −0.50% / −0.27% step time against a 0.04–0.08% within-arm spread, i.e. parity. This is a correctness fix, not a speedup. Register usage goes 57 → 78 VGPRs with zero spills, and the LDS footprint is unchanged at 27392 B.
Per review feedback the refinement was cut from up to ten input passes to at most four (one to seed the first histogram, then one per round, each fusing the previous round's emit). Kernel time on MI355X, 200 iterations after warmup:
randn(no overflow, hot path)randn(no overflow, hot path)narrow_bin(every row refines)narrow_bin(every row refines)two_values(worst case, 32 bits)two_values(worst case, 32 bits)The hot path is untouched; only rows that actually overflow reach the refinement.
Known gap
The CUDA-only
TopKClusterpath still truncates; refining there needs cluster-wide rather than block-local histogram and emit counters, so it is left for a follow-up. #35257 quotes only the register and streaming sites, so that third one is unreported. Related but distinct, #36807 covers the same class of truncation in the AOTtopk.cu.Checklist
Known gap
The CUDA-only
TopKClusterpath still truncates. Refining there is a different routine: the candidate set is split acrosskClusterSizeranks, so each radix pass needs a cluster-wide histogram all-reduce and a cluster-wide emit counter rather than the block-local ones used here. Left for a follow-up rather than grown into this PR — happy to open a tracking issue.CI States
Latest PR Test (Base): 🚫 Run #34665472810
Latest PR Test (Extra): ❌ Run #34665472708
Latest PR Test (AMD ROCm 10): ❌ Run #34665472796