[DSA] Fix IMA in fused top-k v2: write all output slots on tie overflow (#30512)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jimmy Shong
2026-07-08 19:49:26 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 3d96bb9721
commit 4c5fe42be4
@@ -209,6 +209,15 @@ struct TopKConfig {
if (num_ties <= topk) {
if (tx < num_ties) problem.emit(base + tx, tie_buffer[tx].idx);
// Fewer tie candidates than remaining slots (ties beyond kMaxNumTie are
// dropped at collect): pad [num_ties, topk) with -1 ("no token"). The
// transform pass reads all `topk` output slots, and any slot left
// unwritten holds uninitialized staging memory whose page-table
// translation yields a garbage KV index (-> illegal memory access in
// the downstream sparse attention kernel).
for (uint32_t t = num_ties + tx; t < topk; t += kBlockSize) {
problem.emit(base + t, -1u);
}
} else if (num_ties <= kWarpSize) {
if (lane_id >= num_ties || warp_id >= num_ties) return; // some threads are idle
/// NOTE: use long long to avoid mask overflow when num_tie == 32