modelscope / modelscope/mcore-bridge

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

Open
#188 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.