[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:
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
|
||||
|
||||
Reference in New Issue
Block a user