Merged
feat(triton): add chunk_gated_delta_rule_fwd_kkt_solve #135
gcw_K5CvmS79Add on Aug 25
feat(triton): add chunk_gated_delta_rule_fwd_kkt_solve #135
Merged
Pull Request has been successfully merged by ascend-robot.
(Thank you for gcw_K5CvmS79's contribution.)Ggcw_K5CvmS79
Aug 25 Associated issue: [Feature][Q3社区任务]triton算子(chunk_gated_delta_rule_fwd_kkt_solve_kernel)迁移mindspeed-ops仓
Aug 25 Associated issue: [Feature][Q3社区任务]triton算子(chunk_gated_delta_rule_fwd_kkt_solve_kernel)迁移mindspeed-ops仓
atomgit-bot
Aug 25 Comments:
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 描述中说明。


Incorrect?
atomgit-bot
Aug 25 Comments:
Aug 25 Comments:
Aug 25 Added label: stat/needs-squash
Aug 25 Added label: ascend-cla/no
437 messages are collapsed here。 See more
12 days ago Added label: docs-ci-pipeline-success
12 days ago Removed label: ci-pipeline-running
12 days ago Added label: ci-pipeline-passed
12 days ago Closed associated issue
12 days ago Merged pull request
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 wrappermindspeed_ops/api/triton/chunk_gated_delta_rule_fwd_kkt_solve.py:公开 APItests/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 仓内两算子调用):
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。


复现命令:
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、定长输入;其余自动回退未融合链。