QSA Triton kernel fails to compile on Ascend 910B: UB/Cc overflow with default tile sizes

[GAO-Xingyi]

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

Train with default tile size
rank43]: [ERROR] Failed to run BiShengIR pipeline
[rank43]: scf.if %26 {
[rank43]: %reinterpret_cast_63 = memref.reinterpret_cast %arg14 to offset: [%29], sizes: [64], strides: [1] : memref> to memref<64xf32, strided<[1], offset: ?>, #hivm.address_space>
[rank43]: } {limit_sub_block_id0}
[rank43]: hivm.hir.store ins(%subview_64 : memref, #hivm.address_space>) outs(%subview_65 : memref, #hivm.address_space>)
[rank43]: [NOTE] Due to Ub overflow, automatically set enable-code-motion to false, set enable-auto-multi-buffer to false; compilation still failed.
[rank43]: hivm.hir.sync_block_wait[, <PIPE_MTE2>, <PIPE_S>] flag = 3
[rank43]: } {autoblockify.subloop}
[rank43]: loc("/ng-mix02/cogbrain/permanent/xygao19/CODES/TrainSpace/swift_qwen4/mcore-bridge-1.6.4/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py":49:0): error: Failed to run BiShengHIR pipeline
[rank43]: }
[rank43]: %subview_65 = memref.subview %reinterpret_cast_63[0] [%14] [1] : memref<64xf32, strided<[1], offset: ?>, #hivm.address_space> to memref, #hivm.address_space>
[rank43]: loc("/ng-mix02/cogbrain/permanent/xygao19/CODES/TrainSpace/swift_qwen4/mcore-bridge-1.6.4/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py":49:0): error: ub overflow, requires 1918976 bits while 1572864 bits available! (possible reason: tiling basic block is too large or block number is more than what user expect due to multi-buffer feature is enabled and some ops need extra local buffer.)
[rank43]: ///------------------[ERROR][Triton][END]------------------

Error:
Forward kernel compiles with BQ=32 (reduced from 64), but backward (_qsa_bs_dkdv_kernel)
fails with:

• cc overflow, requires 111412 bits while 1048576 bits available
• memref.collapse_shape dim mismatch (32 vs 16)

How to Reproduce / 如何复现

• mcore-bridge: 1.6.4
• CANN: 9.0.0
• triton-ascend: 3.2.2
• megatron-LM: 0.16.0
• MindSpeed: 0.16.0

Additional Information / 补充信息

No response

[hjh0119]

Thanks for the report.

1. We've only tested this kernel on NVIDIA GPUs so far; NPU (triton-ascend) compatibility hasn't been verified, so this issue is indeed possible.
2. We'll add a switch to disable the sparse kernel soon. Note that once the kernel is disabled, QSA falls back to full attention when CP>1 or padding_free is enabled (not an equivalent sparse behavior).
3. If your setup is context_parallel_size 1 with --padding_free false, this kernel is not used at all (the bool-mask path is taken), so you won't be affected.