Closed
[Feature][Q3社区任务]triton算子(chunk_gated_delta_rule_fwd_kkt_solve_kernel)迁移mindspeed-ops仓 #55
且奏长歌Add on  Aug 11Closed at  12 days ago
且奏长歌
且奏长歌Collaborators
Aug 11 Create

任务描述

基于fla-org/flash-linear-attention开源仓库的triton源码,迁移优化triton算子在NPU上跑通及性能优化。

本期任务的具体信息如下:
(1)算子名称:chunk_gated_delta_rule_fwd_kkt_solve_kernel
(2)fla-org/flash-linear-attention开放仓库的算子源码链接:
https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/chunk_fwd.py#L39
(3)迁移指导文档:https://gitcode.com/Ascend/triton-ascend/tree/main/docs/zh/programming_guide
参考 skill 地址:https://gitcode.com/Ascend/agent-skills/tree/master/community/Op/simple-vector-triton-gpu-to-npu
(4)开发平台:Atlas 800T A2或者A3

任务交付件

本期任务为基于fla-org/flash-linear-attention开放仓库代码进行功能扩展,请合入开发代码。
(1)triton算子针对NPU的适配修改优化后的代码
(2)针对该triton算子的实现docu与测试用例UT、ATK用例等
(3)填写昇腾社区开发体验报告

验收标准

(1)基本要求:
功能正常,文档完备,UT、ATK全部通过。
(2)性能要求:
典型场景性能相对两算子性能收益10%以上,并实现精度对齐
(3)精度要求:
对比GPU相同输入算子精度误差小于0.1%;若无GPU标杆,与CPU小算子对齐,精度误差小于0.1%
(4)挑战目标:
A2/A3性能和A100性能比例>0.7
(5)任务完成标准
精度/性能(根据实际要求)达标,PR完成合入,实践文档提交到仓库issue。

补充说明:具体性能对比数据呈现,按 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篇
填写昇腾社区开发体验报告1篇
3)任务完成标准
本次任务完成标准为:
精度/性能(根据实际要求)达标,PR完成合入,实践文档提交到仓库issue。

PR 合入

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

对接人

@LinMingZhe

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

likedislike
且奏长歌且奏长歌Collaborators
Aug 11 Change title to "[Feature]Qwen3.5gdn模块kkt_solve算子性能优化", originally titled "qwen3.5gdn模块kkt_solve算子性能优化".
6 event messages are collapsed here。 See more
YyuhanBaiCollaborators
Aug 11 Modify description of issue
gcw_K5CvmS79Collaborators
Aug 11 Comments:

我认领这个任务

likedislike
且奏长歌且奏长歌Collaborators
Aug 12 Modify description of issue
且奏长歌且奏长歌Collaborators
Aug 12 Change title to "[Feature][Q3社区任务]triton算子(chunk_gated_delta_rule_fwd_kkt_solve_kernel)迁移mindspeed-ops仓", originally titled "[Feature][Q3社区任务] triton算子(chunk_gated_delta_rule_fwd_kkt_solve_kernel)迁移mindspeed-ops仓".
olpk
olpk
Aug 13 Comments:

认领这个任务

likedislike
yuhanBaiCollaborators
Aug 13 Comments:

我认领这个任务

@gcw_K5CvmS79

欢迎认领任务,请参考社区任务池 社区任务池明确该任务的:

完成的截止日期
开发进展反馈
微信答疑群
任务交付注意事项
等信息。如果您同时认领了多项任务,但无法都能进行投入,可以在部分任务中回复退出.

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

likedislike
gcw_K5CvmS79Collaborators
Aug 17 Comments:

正在准备设计文档

likedislike
Ggcw_K5CvmS79Collaborators
Aug 25 Associated pull request
gcw_K5CvmS79Collaborators
Aug 25 Comments:
gcw_K5CvmS79Collaborators
Sep 2 Comments:
XxmzCollaborators
Sep 23 Associated milestone: MindSpeed 26.3.0
SeaAsters
SeaAstersCollaborators
12 days ago Comments:

@LinMingZhe @WendongPang 已提交PR,请检视任务,feat(triton): add chunk_gated_delta_rule_fwd_kkt_solve

@gcw_K5CvmS79

开发者你好,请在任务广场中同步任务进展。

likedislike
ascend-robotascend-robotCollaborators
12 days ago Close issue
ascend-robotascend-robotCollaborators
12 days ago changed custom state from TODO to CLOSED
ascend-robotascend-robotCollaborators
11 days ago Added label: resolved
gcw_K5CvmS79Collaborators
7 days ago Comments:

chunk_gated_delta_rule_fwd_kkt_solve 融合算子实践文档

一、任务概述

  • 任务目标:将 fla-org/flash-linear-attention 的 chunk_gated_delta_rule_fwd_kkt_solve_kernel 迁移并优化到 MindSpeed-Ops(triton-ascend,arch32),用单次调用等价替代 solve_tril(chunk_scaled_dot_kkt_fwd(...)) 两步路径

  • 环境:910B / Atlas 800T A2,bf16

二、方案与实现

保留官方 kkt 构造与 16×16 前代两个 kernel,将官方的两级 merge(16→32→64)加输出预清零压缩为单个 merge kernel,消除 32×32 中间张量的 HBM 往返与两次清零开销,直接输出完整 (I+A)^{-1}。

GPU→NPU 适配要点:门控指数由 exp2 改 tl.exp 并 clamp 到 [-50,50];矩阵乘精度固定 ieee(arch32 cube 无 tf32 累加器);g/beta 转置为 [B,H,T] 恢复连续访存;varlen 路径修正全长与序列长错位。

API 使用约束:仅 chunk_size=64、定长输入走融合路径,其余配置自动回退未融合链,行为完全向后兼容。

三、测试与结果

  • 正确性:UT 53 项 NPU 实跑全过(fp32/bf16/fp16 × 9 shape、±g/±beta、varlen、回退路径);与未融合两算子链位级一致(mismatch=0%);与 CPU fp32/fp64 参考 mean_rel ≤ 1.4e-7,远优于 0.1% 要求。ATK accuracy / performance_device 各 60 case 全部 Pass。

  • 性能(msprof 设备侧 kernel 时间,vs 仓内两算子调用):典型 shape 如 (1, 8192, 2, 128) 0.197→0.161 ms(1.22x)、(1, 16384, 8, 128) 0.714→0.611 ms(1.17x);13 个 shape 中 12 个 ≥1.1x,几何平均 ~1.16x,最大极端 shape(H=32 大行跨距,merge 访存达硬件极限)与两算子持平、不退化。

四、结论与收获

按期完成 issue #55 全部要求:精度位级一致,性能几何平均 1.16x 且全场景不退化,PR 已合入社区。主要收获:掌握了开源 Triton 算子向昇腾平台迁移的适配方法(数值函数替换、精度档位选择、访存布局调整),以及"融合消除中间张量 HBM 往返"这一通用优化思路与自动回退的兼容性设计。

likedislike