modelscope / modelscope/mcore-bridge
QSA Triton kernel fails to compile on Ascend 910B: UB/Cc overflow with default tile sizes
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 101
- Forks
- 43
- Avg merge
- 1d 23h
- Merged PRs (30d)
- 23
Description
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
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py at line 49 and reproduce the QSA Triton kernel compilation using the listed mcore-bridge, CANN, triton-ascend, Megatron-LM, and MindSpeed versions on Ascend 910B. Compare the default tile-size failure with the forward kernel at BQ=32, then inspect _qsa_bs_dkdv_kernel and its backward errors. Done means the forward and backward kernels compile without UB/Cc overflow or the memref.collapse_shape dimension mismatch.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100