From 4c5fe42be4c93a77dcf9d1d9a6e1381171fb2f29 Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 8 Jul 2026 04:49:26 -0700 Subject: [PATCH] [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) --- .../include/sgl_kernel/deepseek_v4/topk_impl.cuh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh index 528d8c6e4..ffdf0f916 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh @@ -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