Merged
feat(triton): add chunk_gated_delta_rule_fwd_kkt_solve #135
feat(triton): add chunk_gated_delta_rule_fwd_kkt_solve #135
Merged
gcw_K5CvmS79Add on Aug 25
gcw_K5CvmS79Collaborators
Aug 25

issue: https://gitcode.com/Ascend/MindSpeed-Ops/issues/55

What this PR does / why we need it?

迁移并优化 fla-org/flash-linear-attention 的 chunk_gated_delta_rule_fwd_kkt_solve_kernel 到 MindSpeed-Ops(triton-ascend, arch32)。新增独立融合算子 chunk_gated_delta_rule_fwd_kkt_solve,一条调用等价于 solve_tril(chunk_scaled_dot_kkt_fwd(...)) 两步路径:保留官方 kkt 构造与 16×16 前代两个 kernel,将官方的两级 merge(16→32→64)+ 输出 torch.zeros 预清零压缩为单个merge kernel,消除 32×32 中间张量的 HBM 往返与两次清零开销。

交付件(5 处,纯新增):

  • mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_fwd_kkt_solve.py:自研 merge kernel + arch32 wrapper
  • mindspeed_ops/api/triton/chunk_gated_delta_rule_fwd_kkt_solve.py:公开 API
  • tests/unit_tests/triton/test_chunk_gated_delta_rule_fwd_kkt_solve.py:UT 53 项
  • tests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve/:ATK 三件套(yaml + generate + runner)
  • docs/triton/chunk_gated_delta_rule_fwd_kkt_solve.md:算子文档

性能(910B/Atlas 800T A2, bf16, msprof 设备侧 kernel 时间,vs 仓内两算子调用):

shape (B,T,H,K) 两算子 ms/iter 融合 ms/iter 加速比
(1, 256, 8, 128) 0.1144 0.0982 1.165
(1, 512, 8, 128) 0.1330 0.1104 1.204
(1, 1024, 8, 128) 0.1515 0.1249 1.213
(1, 2048, 8, 128) 0.1865 0.1543 1.208
(2, 2048, 4, 128) 0.1951 0.1656 1.178
(1, 4096, 4, 128) 0.1946 0.1656 1.175
(1, 8192, 2, 128) 0.1970 0.1610 1.224
(1, 8192, 4, 128) 0.2564 0.2247 1.141
(1, 8192, 8, 128) 0.4318 0.3593 1.202
(1, 8192, 16, 128) 0.7139 0.6365 1.122
(1, 16384, 8, 128) 0.7136 0.6113 1.167
(2, 16384, 16, 128) 2.5123 2.2771 1.103
(2, 16384, 32, 128) 4.8985 4.9368 0.992

12/13 shape ≥ 1.1x,几何平均 ~1.16。最大极端 shape (2,16384,32,128) 与两算子持平:H=32 大行跨距下 merge 访存为硬件极限(官方双 merge 同样受限)。

精度: 与未融合两算子链位级一致(13 基准 + 10 泛化 shape 全部 mismatch=0%,含 T 非 64 倍数、B>1、奇数头数);与 CPU fp32/fp64 参考 mean_rel ≤ 1.4e-7,远优于 0.1% 要求。

Does this PR introduce any user-facing change?

新增公开 API(纯新增,不改任何现有行为,完全向后兼容):

from mindspeed_ops.api.triton.chunk_gated_delta_rule_fwd_kkt_solve import chunk_gated_delta_rule_fwd_kkt_solve

A = chunk_gated_delta_rule_fwd_kkt_solve(
    k, g=None, beta=None, cu_seqlens=None, chunk_size=64, output_dtype=None
)  # k: [B,T,H,K] → A: (I+A)^{-1} [B,T,H,64]
  • chunk_size=64(kernel 硬编码 4 个 BC=16 子块)与 cu_seqlens=None(定长)走融合路径;其他配置自动回退未融合链,行为不变

How was this patch tested?

UT(53 项,NPU 实跑全过,golden 为 CPU 小算子基准): fp32/bf16/fp16 × 9 shape(含 B=2、partial last chunk、NT 奇偶边界、长序列)、±g/±beta、varlen(cu_seqlens 拆段交叉验证)、与仓内未融合链逐位交叉验证、chunk_size≠64 回退路径。

ATK: accuracy 60 case Pass(single_bm,CPU golden);performance_device 60 case Pass。
image.png
image.png

复现命令:

cd tests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve
atk case -f chunk_gated_delta_rule_fwd_kkt_solve.yaml -p generate_chunk_gated_delta_rule_fwd_kkt_solve.py
atk node --backend triton --devices 0 node --backend cpu --devices 0 task \
  -c result/chunk_gated_delta_rule_fwd_kkt_solve/json/all_chunk_gated_delta_rule_fwd_kkt_solve.json \
  --task accuracy -p triton_chunk_gated_delta_rule_fwd_kkt_solve.py

性能复现: msprof --application=... 后读 op_summary_*.csv 的 Task Duration(us),对两算子调用与融合 API 分别累计全部 NPU kernel 时间。

使用约束: 融合路径仅 arch32、chunk_size=64、定长输入;其余自动回退未融合链。

likedislike
Pull Request has been successfully merged by ascend-robot.
(Thank you for gcw_K5CvmS79's contribution.)
Ggcw_K5CvmS79Collaborators
Aug 25 Associated issue: [Feature][Q3社区任务]triton算子(chunk_gated_delta_rule_fwd_kkt_solve_kernel)迁移mindspeed-ops仓
atomgit-bot
atomgit-bot
Aug 25 Comments:

变更摘要

本 PR 将 fla-org/flash-linear-attention 的 Triton 融合算子 chunk_gated_delta_rule_fwd_kkt_solve_kernel 迁移并优化到 MindSpeed-Ops(triton-ascend, arch32),在单个 kernel 内完成 chunk 内 K @ K^T(kkt)+ 下三角求解(tril solve),直接输出完整 (I + A)^{-1},作为 solve_tril(chunk_scaled_dot_kkt_fwd(...)) 两步路径的融合等价替代,省去中间 A 的一次 HBM 往返。变更主要聚焦于:新增融合 kernel 本体与 arch32 wrapper、新增公开 API(含 chunk_size/arch35 守卫)、在顶层 chunk_gated_delta_rule_fwd 中就地加入基于 _FUSED_KKT_SOLVE_T_MAX = 1024 的分流逻辑(小/中 T 走融合算子,大 T 与 varlen/arch35/非 64 chunk 走原两算子路径,对既有调用方向后兼容),并配套 UT、ATK 用例与算子文档;diff 中还附带新增了未在 PR 描述中提及的 layer_norm_fwd_1pass kernel 及其 API 与测试。

主要改动

  • 新增融合 kernel 与公开 API:新增 mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_fwd_kkt_solve.py(kernel + arch32 wrapper,tl.make_block_ptr 加载、寄存器内计算 10 个下三角 [BC,BC] 块并做前向替换与块合并)和 mindspeed_ops/api/triton/chunk_gated_delta_rule_fwd_kkt_solve.py(公开 API:chunk_size != 64 抛 ValueError,is_arch35() 抛 NotImplementedError)。
  • 顶层 chunk_gated_delta_rule_fwd 增加分流:在 mindspeed_ops/api/triton/chunk_gated_delta_rule.py 中新增 _FUSED_KKT_SOLVE_T_MAX = 1024 阈值,满足 chunk_size == 64、cu_seqlens is None、非 arch35 且 T <= 1024 时走 chunk_gated_delta_rule_fwd_kkt_solve,否则回退 chunk_scaled_dot_kkt_fwd + solve_tril,保证全场景不退化且输出与原路径等价。
  • GPU→NPU 适配关键改动:kernel 内门控指数由 exp2 改为 tl.exp 并 clamp 到 [-50,50];矩阵乘精度固定为 ieee(SOLVE_TRIL_DOT_PRECISION,arch32 cube 无 tf32 累加器);wrapper 将 g/beta 从 [B,T,H] 转置为 [B,H,T](步长 1)以恢复连续访存;varlen 路径通过 prepare_chunk_indices 与 T_local 修正全长 T 与序列长错位的 bug。
  • 测试与文档配套:新增 tests/unit_tests/triton/test_chunk_gated_delta_rule_fwd_kkt_solve.py(45 项 UT,覆盖 fp32/bf16/fp16、±g/±beta、varlen、与仓内非融合路径交叉验证、非法 chunk_size 报错及 TestTopLevelDispatch 顶层分流两条路径)、tests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve/ 的 ATK 三件套(yaml + generate + runner),以及 docs/triton/ 下算子介绍与开发体验报告两份文档。
  • 附带新增 layer_norm_fwd_1pass:diff 中同时新增了 mindspeed_ops/api/triton/layer_norm_fwd_1pass.py、mindspeed_ops/arch32/triton/layer_norm_fwd_1pass.py 及其 UT/ATK 用例(基于 UB 容量动态计算 BT、1D 物理核网格任务分发、支持 is_rms_norm/ngroups/gating 前后置),arch35 上抛 NotImplementedError,此部分未在 PR 描述中说明。
likedislike
Incorrect?
atomgit-bot
atomgit-bot
Aug 25 Comments:

代码审查

✅ 未发现问题

likedislike
Incorrect?
ascend-robotascend-robotCollaborators
Aug 25 Added label: stat/needs-squash
ascend-robotascend-robotCollaborators
Aug 25 Added label: ascend-cla/no
437 messages are collapsed here。 See more
ascend-robotascend-robotCollaborators
12 days ago Added label: docs-ci-pipeline-success
AtlasAccountAtlasAccountCollaborators
12 days ago Removed label: ci-pipeline-running
AtlasAccountAtlasAccountCollaborators
12 days ago Added label: ci-pipeline-passed
ascend-robotascend-robotCollaborators
12 days ago Closed associated issue
ascend-robotascend-robotCollaborators
12 days ago Merged pull request