已开启

[Feature]: triton算子(mamba_mimo_bwd_fwd_kernel)迁移mindspeed-ops仓 #38

Ronald1995创建于 6月11日

Ronald1995成员

6月11日创建

提交提案之前，请先检索仓库内是否已有相同的提案，如已有请在同一提案中进行讨论。

💻 需求背景、当前现状、期望实现的功能内容、具体的设计方案、以及测试方案

任务描述

基于state-spaces/mamba开源仓库的triton源码，迁移优化triton算子在NPU上跑通及性能优化。
本期任务的具体信息如下：
（1）算子名称：mamba_mimo_bwd_fwd_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%
性能：
对比原triton算子性能提升20%
2）实践文档：
triton算子介绍1篇
3）任务完成标准
本次任务完成标准为：
精度/性能（根据实际要求）达标，PR完成合入，实践文档提交到仓库issue。

PR 合入

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

对接人

LinSHua

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

替代方案

补充说明

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

R Ronald1995成员

6月11日添加了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日评论：

目前进展：

triton 迁移:先把原始算子迁到 triton-ascend,删掉 SEGSUM 中间量、改成片上从 dA_cs 重建，拿到能跑的基线。
把小矩阵乘从 Vector 核搬到 Cube 核：针对 scan(状态递归)部分计算有大量 K=16 的小矩阵乘，尝试把它卸到专门做矩阵乘的 Cube 核 → 40.9us,快了 47×。中间试过 Brcb 向量化外积那条路，目前还没跑通。
希望进一步提高性能，打满硬件计算，但在 triton 遇到瓶颈。尝试实现AscendC算子。
dmimo_o 重构:把 2560 个极小的 M=16 矩阵乘合并成 3 个 M=32 的大矩阵乘,并用自然内存布局消掉转置胶水。

结果: 输出精度与GPU tilelang基线相比误差 <0.1%；目前加速比已经达标，但在继续尝试优化。目前 msprof 实测是 scalar 瓶颈，未达性能上限。
下一步计划：继续优化。

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-fwd-kernel
预计将提交：pytorch小算子实现，triton优化前代码，triton优化后代码，优化记录，测试用例，算子迁移说明。
另外正尝试打包ascendc算子到mindspeed-ops仓库。

L liuzhexu成员

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

bitszh3271成员

7月14日评论：

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

Shape	Speedup
cg_N16_P64_R4_C8	1.81×
cg_N32_P64_R4_C16	1.86×
cg_N64_P64_R4_C16	1.84×
cg_N128_P64_R4_C16	1.89×
cg_N256_P64_R4_C8	1.49×
cg_N64_P128_R4_C16	2.07×
cg_N128_P32_R4_C16	1.93×
cg_N128_P128_R4_C8	1.69×
cg_N128_P64_R8_C8	5.78×
cg_N128_P64_R2_C32	1.12×
cg_N128_P64_R1_C64	1.01×
几何平均	1.83

已满足加速要求。算子已在私仓提交，正在写相关文档。
ascendc 方向，前期agent将bwd_fwd拆成若干个kernel实现，目前正在尝试写CV融合的 megakernel.

bitszh3271成员

7月16日评论：

triton-kernel pr已提交， https://gitcode.com/Ascend/MindSpeed-Ops/pull/101
ascendc kernel 正在优化。

bitszh3271成员

7月24日评论：

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

8月3日评论：

我要认领这个任务

chenchen

8月11日评论：

认领这个任务

bitszh3271成员

8月13日评论：

@Ronald1995

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计算），难以继续优化。

梵高的呐喊

2 天前评论：

认领该任务