已开启

【Feature】triton算子(mamba_mimo_bwd_bwd_kernel)迁移mindspeed-ops仓 #23

LinShua创建于 5月18日

5月18日创建

任务描述

基于state-spaces/mamba开源仓库的triton源码，迁移优化triton算子在NPU上跑通及性能优化。
本期任务的具体信息如下：
（1）算子名称：mamba_mimo_bwd_bwd_kernel
（2）state-spaces/mamba开放仓库的算子源码链接：
https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_bwd.py
（3）迁移指导文档：https://gitcode.com/Ascend/triton-ascend/blob/main/docs/zh/programming_guide.md
参考 skill 地址：https://gitcode.com/Ascend/agent-skills/tree/master/skills/simple-vector-triton-gpu-to-npu
（4）开发平台：Atlas 800T A2或者A3

验收标准

一、任务交付件
本期任务为基于state-spaces/mamba开放仓库代码进行功能扩展，请合入开发代码。主要开发点如下：
（1） triton算子针对NPU的适配修改优化后的代码
（2）针对该triton算子的测试用例UT
二、验收标准：
1）精度/性能要求
精度：
对比GPU相同输入算子精度误差小于0.1%
若无GPU标杆，与CPU小算子对齐，精度误差小于0.1%
性能：
补充说明：具体性能对比数据呈现，按 CV类算子达到0.7x竞品， VV类算子达到0.9x竞品
以下以竞品A100为例说明计算具体的性能对比要求:

1.先计算当前运行NPU设备与对应竞品GPU的算力换算比例

硬件规格	910B2	A100	910B2/A100
Cube算力	354	312	1.13x
Vector算力	11.06	19.5	0.56x
带宽	1800	2039	0.88x

2.根据算力换算比例计算对应算子类在不同bound场景的实际性能要求

910B2 VS A100

bound场景	CV类算子(0.7X)	VV类算子(0.9x)
Cube算力bound场景	0.8x	1.0X
Vector算力bound场景	0.4x	0.5X
Memory bound场景（HBM）	0.6x	0.8x

计算式子：以CV类算子的Cube算力bound场景为例：1.13 * 0.7 =0.8
2）实践文档：
triton算子介绍1篇
3）任务完成标准
本次任务完成标准为：
精度/性能（根据实际要求）达标，PR完成合入，实践文档提交到仓库issue。

PR 合入

本地完成测试验证后，向MindSpeed-Ops的master分支发起PR。

对接人

LinSHua

欢迎加入社区，感谢您对社区的贡献 🎉!

L LinShua成员

6月2日添加了label：feature

Ronald1995成员

6月11日评论：

/label add good-first-issue

ascend-robot成员

6月11日添加了label：good-first-issue

wangx700成员

6月12日评论：

/label add triage-review

ascend-robot成员

6月12日添加了label：triage-review

Tsuki成员

6月25日评论：

认领这个任务

Ronald1995成员

6月25日评论：

认领这个任务

@xiaomaomao666

欢迎认领任务，请参考https://gitcode.com/Ascend/MindSpeed-Ops/issues/3 社区任务池明确该任务的：

完成的截止日期
开发进展反馈
微信答疑群
任务交付注意事项

等信息。如果您同时认领了多项任务，但无法都能进行投入，可以在部分任务中回复退出.

麻烦您加入到对应微信群，群备注名修改为"社区任务+您的gitcode账号", 后续有相关消息和问题都可以在微信群咨询答疑。等您加入到微信群后，我这边会在社区任务池里面登记任务责任人。

Tsuki成员

6月26日评论：

退出

bitszh3271成员

6月28日评论：

认领这个任务

gcw_MTeIQk9o

6月29日评论：

认领这个任务

yuhanBai成员

6月30日评论：

认领这个任务

@bitszh3271

你好，感谢您对昇腾平台的支持，请及时在本issue评论区更新进展，超过一周无进展更新自动视为放弃该任务。：）

yuhanBai成员

6月30日评论：

认领这个任务

@gcw_MTeIQk9o

你好，感谢您对昇腾平台的支持，当前一个任务只分配一个开发，根据评论认领时间确定，您可以后续观察该任务是否被释放退出，或者认领其他任务 https://gitcode.com/org/Ascend/discussions/4

zeshengzong成员

7月2日关联了看板：@zeshengzong的看板 20260702

ascend-robot成员

7月3日关联了看板：MindStudio ISSUE管理

gcw_K5CvmS79成员

7月3日评论：

认领这个任务

bitszh3271成员

7月4日评论：

当前工作:

处理迁移精度问题：原迁移版 12 个梯度里有 2 个过不了,通过 DGAMMA 累加顺序重排(避开小量放大)、删 SEGSUM 等，修到全过；与GPU Tilelang实现对比，精度达标。
triton 层优化：把 scan 内只依赖 dt/trap 的离散化逐元素链(exp/sigmoid)整体挪到并行 pre-kernel,从串行关键路径上移走 → 5.4–5.9×;再合并 scratch DMA(52→8)叠加到 6.7×。
实现 AscendC 原生算子:同 bwd_fwd优化思路，将 9 个小矩阵乘全卸 Cube; 设计 two-pass 解耦——先算 chunk 内的量,反向递归用纯加乘。

结果：
12 个梯度全部 <1e-3; 按仓库标准口径 wall-clock 6.60× / 纯 device 9.26× / 持续吞吐 6.78×。msprof 实测是 host/launch 瓶颈；device 侧 Cube 已喂饱(cube% 87)，但 Vector 的 scalar 仍主导。

下一步：
尝试将多个小asc算子融合成大算子。

yuhanBai成员

7月7日评论：

认领这个任务

@gcw_K5CvmS79

你好，感谢您对昇腾平台的支持，当前一个任务只分配一个开发，根据评论认领时间确定。请关注https://gitcode.com/Ascend/MindSpeed-Ops/issues/3任务是否释放或者认领其他任务 https://gitcode.com/org/Ascend/discussions/4

bitszh3271成员

7月9日评论：

triton版本算子已经准备就绪，待完善周边资料提交pr：https://gitcode.com/bitszh3271/MindSpeed-Ops/tree/feat/mamba3-mimo-bwd-bwd-kernel。
预计将提交：pytorch小算子实现，triton优化前代码，triton优化后代码，优化记录，测试用例，算子迁移说明。
triton最终优化结果：
device-sum 242.6us，scan 152.55us，prepass 46.68us，fuse_post 35.81us，单次 ReduceSum 7.52us。相比 C16 fuse 640.8us、C32 fuse 407.5us，C64 split fuse 是当前最优。
另外正尝试打包ascendc算子到mindspeed-ops仓库。

L liuzhexu成员

7月11日关联了里程碑：MindSpeed 26.2.0

bitszh3271成员

7月14日评论：

花了几天时间对齐补全官方测试用例的全部91 testcase。triton 算子版本优化结果如下，使用 mamba-3 官方 11shape 测试用例。

Shape	#23 Speedup	备注
cg_N16_P64_R4_C8	13.52×	(aux-fold 前 6.40×)
cg_N32_P64_R4_C16	29.45×	(第一版优化时 17.28×)
cg_N64_P64_R4_C16	31.55×	(第一版优化时 18.90×)
cg_N128_P64_R4_C16	34.81×	(第一版优化时 21.64×)
cg_N256_P64_R4_C8	—	stage0 基线 UB 溢出(无基线)
cg_N64_P128_R4_C16	29.62×	(第一版优化时 18.51×)
cg_N128_P32_R4_C16	36.05×	(第一版优化时 21.74×)
cg_N128_P128_R4_C8	—	stage0 基线 UB 溢出(无基线)
cg_N128_P64_R8_C8	8.16×	(R8 不走 aux-fold)
cg_N128_P64_R2_C32	—	stage0 基线 UB 溢出(无基线)
cg_N128_P64_R1_C64	—	stage0 基线 UB 溢出(无基线)

baseline移植（直接迁移）在某些shape下会爆UB，因此无加速比。
主要优化点为scan/aux R×R 微矩阵 scalar 墙，将其合并为大矩阵计算，此外算子融合，调整分块也是主要优化项。
算子已在私仓提交，正在写相关文档。
ascendc 方向，前期agent将bwd_fwd拆成若干个kernel实现，目前正在尝试写CV融合的 megakernel.

bitszh3271成员

7月16日评论：

triton-kernel已经提交：https://gitcode.com/Ascend/MindSpeed-Ops/pull/100
ascendc-kernel正在优化

bitszh3271成员

7月24日评论：

ascendc kernel实现见fork仓库，目前asc kernel约1~2倍triton性能。考虑到合入仓库较为麻烦，暂时推迟。

8月3日评论：

我要认领这个任务

L LinShua成员

8月7日修改了issue 的描述

bitszh3271成员

8月13日评论：

@LinShua

AscendC Official 11 逐 Shape 三轮中位数统计（单位：ms）

Ratio 定义为 A100耗时 / 910B3耗时。

Official Shape	mimo bwdfwd A100	mimo bwdfwd AscendC	mimo bwdfwd Ratio	mimo bwdbwd A100	mimo bwdbwd AscendC v119c	mimo bwdbwd Ratio
N16_P64_R4_C8_BB128	1.538	7.842	0.1962x	1.940	15.498	0.1252x
N32_P64_R4_C16_BB256	1.479	5.536	0.2671x	2.259	10.252	0.2203x
N64_P64_R4_C16_BB256	1.682	6.306	0.2668x	2.669	11.132	0.2398x
N128_P64_R4_C16_BB256	2.323	7.981	0.2910x	3.935	14.013	0.2808x
N256_P64_R4_C8_BB256	5.181	14.565	0.3557x	6.950	26.396	0.2633x
N64_P128_R4_C16_BB256	2.328	7.987	0.2914x	3.468	13.123	0.2643x
N128_P32_R4_C16_BB256	1.740	7.046	0.2469x	3.771	12.630	0.2986x
N128_P128_R4_C8_BB256	5.009	13.247	0.3781x	5.913	22.522	0.2625x
N128_P64_R8_C8_BB256	4.579	16.257	0.2817x	8.229	37.347	0.2203x
N128_P64_R2_C32_BB256	1.269	4.193	0.3025x	1.994	8.910	0.2238x
N128_P64_R1_C64_BB256	0.773	2.330	0.3320x	1.294	4.122	0.3139x
11-shape GM	2.135	7.423	0.2876x	3.311	13.761	0.2406x

Triton Official 11 逐 Shape 三轮中位数统计（单位：ms）

Ratio 定义为 A100耗时 / 910B3耗时。

Official Shape	mimo bwdfwd A100	mimo bwdfwd AscendC	mimo bwdfwd Ratio	mimo bwdbwd A100	mimo bwdbwd AscendC v119c	mimo bwdbwd Ratio
N16_P64_R4_C8_BB128	1.538	30.424	0.0506x	1.940	30.662	0.0633x
N32_P64_R4_C16_BB256	1.479	15.124	0.0978x	2.259	18.876	0.1197x
N64_P64_R4_C16_BB256	1.682	14.772	0.1139x	2.669	20.209	0.1321x
N128_P64_R4_C16_BB256	2.323	15.461	0.1502x	3.935	24.499	0.1606x
N256_P64_R4_C8_BB256	5.181	42.057	0.1232x	6.950	45.009	0.1544x
N64_P128_R4_C16_BB256	2.328	18.564	0.1254x	3.468	23.432	0.1480x
N128_P32_R4_C16_BB256	1.740	14.540	0.1197x	3.771	22.523	0.1674x
N128_P128_R4_C8_BB256	5.009	41.127	0.1218x	5.913	38.051	0.1554x
N128_P64_R8_C8_BB256	4.579	45.983	0.0996x	8.229	50.917	0.1616x
N128_P64_R2_C32_BB256	1.269	7.345	0.1727x	1.994	16.912	0.1179x
N128_P64_R1_C64_BB256	0.773	2.569	0.3011x	1.294	4.231	0.3058x
11-shape GM	2.135	17.338	0.1231x	3.311	22.980	0.1441x

反向传播算子复杂，输入输出多，主要受限UB大小，HBM搬运效率，以及算法中包含很多大batch小矩阵乘法（MIMO计算），难以继续优化。

梵高的呐喊

9月17日评论：

认领该任务

X xmz成员

9月23日关联了里程碑：MindSpeed 26.3.0