已开启
【Feature】triton算子(mamba3_siso_bwd_kernel_dqkv)迁移mindspeed-ops仓 #15
LinShua创建于  5月18日
LinShua成员
5月18日 创建

任务描述

基于state-spaces/mamba开源仓库的triton源码,迁移优化triton算子在NPU上跑通及性能优化。
本期任务的具体信息如下:
(1)算子名称:mamba3_siso_bwd_kernel_dqkv
(2)state-spaces/mamba开放仓库的算子源码链接:
https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/triton/mamba3/mamba3_siso_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

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

likedislike
LLinShua成员
6月4日 添加了label:feature
Ronald1995成员
6月11日 评论:

/label add good-first-issue

likedislike
ascend-robotascend-robot成员
6月11日 添加了label:good-first-issue
wangx700
wangx700成员
6月12日 评论:

/label add triage-review

likedislike
ascend-robotascend-robot成员
6月12日 添加了label:triage-review
redlightASl
6月26日 评论:

认领这个任务

likedislike
~~灏
~~灏
6月26日 评论:

认领这个任务

likedislike
bitszh3271成员
6月27日 评论:

认领这个任务

likedislike
gcw_MTeIQk9o
6月29日 评论:

认领这个任务

likedislike
Ronald1995成员
6月29日 评论:

认领这个任务

@redlightasl1

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

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

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

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

likedislike
Ronald1995成员
6月29日 评论:

认领这个任务

@gcw_MTeIQk9o

@bitszh3271

@Morikunn

欢迎参加认领社区任务,当前一个任务只分配一个开发,根据评论认领时间确定,您可以后续观察该任务是否被释放退出,或者认领其他任务 https://gitcode.com/org/Ascend/discussions/4

likedislike
yuhanBai成员
6月30日 评论:

认领这个任务

@redlightasl1

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

likedislike
RredlightASl
7月1日 关联了pull request:feat(triton): add mamba3 siso bwd dqkv kernel
redlightASl
7月1日 评论:

已提交 MR: https://gitcode.com/Ascend/MindSpeed-Ops/merge_requests/86

本次迁移内容:

  • 新增 mamba3_siso_bwd_kernel_dqkv 的 arch32 Triton-Ascend 生产实现和公开 API。
  • 增加 ATK/pytest 测试、自包含 self-check 脚本和中文说明文档。
  • 文档中已说明原始 upstream Triton 版本的 split/fused 组成、运行时路径选择,以及本 PR 的 single-chunk fused、small fused、V/K/KV tiled fused 和同 backend split 调度。
  • 速度 baseline 使用 PR 内 split Triton-Ascend baseline 现场 profiler 计时;torch reference 仅用于精度校验。

本地验证摘要:

  • full 自证报告已写入 docs/triton/mamba3_siso_bwd_kernel_dqkv.md:20/20 case PASS,Ascend AVA 20/20 case >= 1.2x,geomean 5.318358x,min 1.51935x。
  • python -m py_compile ...:PASS。
  • 外部 OpForge 路径/工件扫描:无命中。
  • disallowed timing source 扫描:无命中。
  • git diff --check origin/master...HEAD:PASS。
  • 非 0 号 NPU quick no-timing 自证:PASS(2/2 代表 case,未重复跑 full 性能)。
likedislike
zeshengzongzeshengzong成员
7月2日 关联了看板:@zeshengzong的看板 20260702
ascend-robotascend-robot成员
7月3日 关联了看板:MindStudio ISSUE管理
Lliuzhexu成员
7月11日 关联了里程碑:MindSpeed 26.2.0
8月3日 评论:

我要认领这个任务

likedislike
LLinShua成员
8月7日 修改了issue 的描述
梵高的呐喊
梵高的呐喊
7 天前 评论:

认领该任务

likedislike
quipiter
7 天前 评论:

认领这个任务

likedislike
Xxmz成员
1 天前 关联了里程碑:MindSpeed 26.3.0