feat: add Triton-Ascend implementation of Mamba-3 mamba_mimo_bwd

已开启

feat: add Triton-Ascend implementation of Mamba-3 mamba_mimo_bwd_fwd #125

hahahascend创建于 8月19日从hahahascend/MindSpeed-Ops： codex/mamba-mimo-bwd-fwd-triton 合入到 Ascend/MindSpeed-Ops： master

feat: add Triton-Ascend implementation of Mamba-3 mamba_mimo_bwd_fwd #125

已开启

hahahascend创建于 8月19日

从hahahascend/MindSpeed-Ops codex/mamba-mimo-bwd-fwd-triton合入到Ascend/MindSpeed-Ops master

hahahascend

8月19日

关联 Issue

#22
https://gitcode.com/Ascend/MindSpeed-Ops/issues/22

What this PR does / why we need it?

基于 state-spaces/mamba 提交 48582f0，新增 Mamba-3 mamba_mimo_bwd_fwd 算子的 Triton-Ascend arch32 实现。

主要修改：

新增 arch32 Triton kernel、参数校验和公开 Python API；
对齐上游功能分支，支持可选 D/Z、两种 DOUT 布局、GQA 和非整除尾 chunk；
优化 control 预处理、Q/K rotary 和 recurrent state scan，减少小 kernel 调度和冗余计算；
新增 CPU FP32 golden、13 项 UT、ATK 配置和算子实践文档。

Does this PR introduce any user-facing change?

是。本 PR 新增以下 Triton API：

from mindspeed_ops.api.triton.mamba_mimo_bwd_fwd import mamba_mimo_bwd_fwd

当前实现支持 Ascend arch32 和 BF16 输入。详细输入约束和使用方法见 docs/triton/mamba_mimo_bwd_fwd.md。

How was this patch tested?

1. 测试对象和统一负载

NPU 和 H100 均只测试 Mamba-3 第一阶段反向算子 mamba_mimo_bwd_fwd，不包含第二阶段 backward kernel。

性能负载：

B=1, S=8192, H=16, G=1, N=P=64, R=4, C=16
hasZ=True, hasD=True, reduceO=True
BF16 inputs

统一计时边界：

输入创建、文件读取和 H2D 不计时；
Triton/TileLang 编译不计时；
每个独立进程预热 10 次；
每个进程采集 30 次同步端到端 API 时间；
共运行 3 个独立进程，正式统计包含 90 个样本；
每次调用均包含五个输出和内部 workspace 的分配；
不包含第二阶段 backward kernel；
预分配 workspace 的结果只用于诊断，不参与验收。

2. 精度测试

精度以独立 CPU FP32 实现作为 golden，分别比较五个输出：

DMIMO_O
STATES
DZ
DMIMO_Z
QK_DOT

精度 shape：

B=1, S=32, H=1, G=1, N=P=64, R=4, C=16
hasZ=True, hasD=True, reduceO=True

精度结果：

输出	Error ratio	Max absolute error
`DMIMO_O`	`1.02081e-6`	`1.43051e-5`
`STATES`	`1.52408e-6`	`3.05176e-5`
`DZ`	`7.89688e-5`	`4.88281e-4`
`DMIMO_Z`	`1.07553e-6`	`1.62125e-5`
`QK_DOT`	`0`	`0`

五个输出的 error ratio 均小于验收阈值 0.001，精度达标。

完整 UT 还覆盖：

D/Z/reduceO 四种功能组合；
GQA；
非整除尾 chunk；
多组 (N, P, R, C) 参数；
R=1 fallback；
非法 shape 和 dtype 校验。

UT 结果：

13 passed

3. H100 竞品数据

竞品环境：

GPU: NVIDIA H100 80GB HBM3 SXM
功耗上限: 700 W
MIG: disabled
状态: P0，测试期间频率保持 1980 MHz
Mamba commit: 48582f0592f1d55b441cd963aa5f9226b083f02a

用于折算的 H100 硬件规格：

指标	H100 SXM
Dense BF16 Tensor 算力	`989.5 TFLOPS`
FP32 Vector 算力	`67 TFLOPS`
HBM 带宽	`3350 GB/s`

未使用 sparse BF16 算力。

H100 三个独立进程结果：

process medians: 8.236963 / 8.201160 / 8.188458 ms
all 90 median:   8.201160 ms
P20 / P80:       8.144154 / 8.264276 ms
min / max:       8.090053 / 8.697119 ms
sample stddev:   0.102551 ms

预分配输出的诊断中位数为 8.075561 ms，由于不包含正式 API 中的分配开销，因此不用于验收换算。

4. NPU 性能数据

NPU 环境：

Device: Ascend910_9382
CANN: 8.5

NPU 三个独立进程结果：

process medians: 59.604352 / 58.998696 / 59.782081 ms
all 90 median:   59.592497 ms
P20 / P80:       59.045436 / 59.806263 ms
min / max:       58.541977 / 60.270459 ms
sample stddev:   0.424671 ms

5. 性能分类和折算方法

该实现同时包含矩阵乘和大量 Vector 计算，因此按照验收标准归为 CV 类算子。

NPU profile 显示主要时间消耗在 control、Q/K bias/rotary、PsiV、state update、output 和 QK gather 等阶段。这些阶段包含较多 exp/sigmoid/sin/cos、逐元素乘加、状态更新、数据重排和小矩阵计算，性能不由大规模 Cube GEMM 单独主导。因此采用 CV 类 Vector-bound 口径进行验收。

验收使用以下硬件 Vector 算力：

设备	FP32 Vector 算力
Ascend 910B2	`11.06 TFLOPS`
NVIDIA H100 SXM	`67 TFLOPS`

先计算 NPU 与竞品的 Vector 算力比例：

910B2 / H100 Vector ratio
= 11.06 / 67
= 0.1650746

CV 类算子的验收系数为 0.7，因此要求的实际性能比为：

required_ratio
= 11.06 / 67 * 0.7
= 0.115552

这里的性能比定义为：

T_H100 / T_NPU

由于时间越小代表性能越高，因此要求：

T_H100 / T_NPU >= 0.115552

也可以换算为 NPU 允许的最大执行时间：

NPU time limit
= T_H100 / required_ratio
= 8.201160 / 0.115552
= 70.973761 ms

6. 达标结论

NPU 实测中位数为：

T_NPU = 59.592497 ms

实测性能比：

measured_ratio
= T_H100 / T_NPU
= 8.201160 / 59.592497
= 0.137621

与要求比较：

measured_ratio = 0.137621
required_ratio = 0.115552

0.137621 > 0.115552

因此性能达标。

验收余量：

性能比高于最低要求约 19.10%
NPU 时间低于 70.973761 ms 上限约 16.04%

7. 其他检查

UT:         13 passed
Ruff:       passed
Pylint:     passed
codespell:  passed
diff check: passed

ATK 精度配置、输入生成脚本和算子入口已提交到：

tests/atk_tests/triton/mamba_mimo_bwd_fwd/

当前开发环境未安装 ATK Python 包，因此本地未生成 ATK 报告；本 PR 不将 ATK 配置存在等同于 ATK 已执行通过。

合并受阻

H hahahascend

8月19日关联了issue：【Feature】triton算子(_layer_norm_bwd_kernel)迁移mindspeed-ops仓

atomgit-bot

8月19日评论：

变更摘要

本 PR 基于 state-spaces/mamba 提交 48582f0，新增 Mamba-3 第一阶段反向算子 mamba_mimo_bwd_fwd 的 Triton-Ascend arch32 实现，并配套公开 Python API、参数校验、CPU FP32 golden、13 项单元测试、ATK 配置及算子文档。实现支持可选 D/Z、两种 DOUT 布局（mimo_o 是否传入决定是否 reduce）、GQA 和非整除尾 chunk，内部以 FP32 中间量计算并保持与上游一致的输出 dtype；性能验收按 CV 类 Vector-bound 口径折算，NPU（Ascend910_9382）实测中位数 59.59 ms，相对 H100 竞品（8.20 ms）的性能比 0.137621 高于要求值 0.115552，且五个输出（DMIMO_O、STATES、DZ、DMIMO_Z、QK_DOT）的 error ratio 均小于阈值 0.001，UT 结果为 13 passed。

主要改动

新增公开 Triton API：mindspeed_ops/api/triton/mamba_mimo_bwd_fwd.py 定义 MambaMimoBwdFwdOutputs NamedTuple（含 dmimo_o、states、dz、dmimo_z、qk_dot）与入口函数 mamba_mimo_bwd_fwd，并在 mindspeed_ops/api/triton/__init__.py 中导出；_validate_inputs 负责校验 arch32 限制（is_arch35() 抛 NotImplementedError）、BF16 dtype、NPU device、chunk_size * R 范围、rotary_dim_divisor 取值、mimo_o/z/mimo_z 的组合约束及全部输入 shape/dtype。
新增 arch32 Triton kernel 实现：mindspeed_ops/arch32/triton/mamba_mimo_bwd_fwd.py 以多个 @triton.jit kernel 实现完整计算流程，包括 control 预处理（mamba_mimo_token_control_kernel、mamba_mimo_rotary_control_kernel、mamba_mimo_causal_control_kernel）、Q/K bias 与 rotary 处理（分组 tile 路径 mamba_mimo_qk_bias_kernel + mamba_mimo_qk_rotary_patch_kernel，大 tile 平铺 fallback mamba_mimo_preprocess_kernel）、mamba_mimo_psi_kernel、持久化 recurrent state scan（mamba_mimo_state_scan_kernel）、mamba_mimo_output_kernel、mamba_mimo_qk_gather_kernel、mamba_mimo_projection_grads_kernel 与 mamba_mimo_dz_kernel，中间结果以 FP32 workspace 承载，输出再转回输入 dtype。
新增 CPU FP32 golden 参考实现：tests/unit_tests/triton/mamba_mimo_bwd_fwd_reference.py 提供 mamba_mimo_bwd_fwd_golden，覆盖 chunk padding（_pad）、rotary（_rotate）、因果 mask、state scan、可选 D/Z 与 reduced/non-reduced DOUT 两种分支，作为精度比对基准。
新增 13 项单元测试：tests/unit_tests/triton/test_mamba_mimo_bwd_fwd.py 通过 _run_case 以 get_err_ratio 校验五个输出，覆盖 D/Z/reduceO 四种功能组合、GQA（heads=2）、非整除尾 chunk（seqlen=20）、多组 (N, P, R, C) 参数（含 R=1 fallback）以及非法 shape/dtype 校验。
新增 ATK 配置并调整代码检查：tests/atk_tests/triton/mamba_mimo_bwd_fwd/ 新增 case 生成器 generate_mamba_mimo_bwd_fwd.py（内置 4 组用例）、mamba_mimo_bwd_fwd.yaml（bf16_error: 0.001）及注册 torch golden 与 Triton API 的 triton_mamba_mimo_bwd_fwd.py；.pre-commit-config.yaml 将 DOUT,dout 加入 codespell 忽略词表。

不准确？

atomgit-bot

8月19日评论：

代码审查

I have completed my review within budget. Let me finalize with the closing summary.

Closing Summary

已审查文件清单

1. .pre-commit-config.yaml — no issues
仅向 codespell 忽略列表追加 DOUT,dout 两个词，属于正常的拼写白名单维护，无正确性问题。

2. mindspeed_ops/api/triton/__init__.py — no issues
新增包 __init__.py，从 mamba_mimo_bwd_fwd 导出两个公开符号。导入链无循环依赖（mamba_mimo_bwd_fwd 仅在函数体内惰性导入 arch32 kernel，顶层只依赖 torch 与 mindspeed_ops.utils，而后者不反向依赖本包）。

3. mindspeed_ops/api/triton/mamba_mimo_bwd_fwd.py — 2 个 P3 发现

P3（conf 0.5）：_validate_inputs 只校验 seqlen/rank/groups 为正，漏掉了 batch 维度的正值校验，batch=0 会静默通过校验并以空 grid 启动 kernel、返回空输出，掩盖参数错误。
P3（conf 0.4）：is_arch35() 基于进程默认设备（torch_npu.npu.get_device_name()）而非实际张量所在 q.device 判断架构，混合架构多卡环境下可能误抛 NotImplementedError 或误放行；且该函数被 @lru_cache 缓存，进程内切换默认设备后结果不更新。

统计与总体风险判断

共 2 个发现，均为 P3 低严重度（边界校验缺口与架构门控的设备绑定问题），无 P0/P1/P2 问题。
总体风险：低。核心校验逻辑（shape/dtype/device、GQA 整除、chunk 边界、reduceO/Z 组合一致性）与 arch32 kernel 的索引约定逐项核对一致，MambaMimoBwdFwdOutputs 字段顺序与 kernel 返回的 5 元组也完全匹配；两个发现均只在边界输入或多设备混合架构场景下触发，不构成常规路径的功能回归。

我已完成对所有 7 个变更文件的审查。核心发现已上报，现给出总结。

审查总结

各文件审查结果

mindspeed_ops/arch32/triton/mamba_mimo_bwd_fwd.py — 发现 1 个问题（P2）：
- raw_output 越界写竞争（第 455–463 行）：mamba_mimo_output_kernel 用 (BF, BP) 块写入 (S*R, P) 平铺缓冲，当 fused = chunk_size*rank 非 2 的幂时（BF > F），非末尾 chunk 的 store 会把零值写进相邻 chunk 的 raw_output 区域，与该 chunk 的 program 并发竞争，导致 DMIMO_O/DZ 非确定性损坏。所有现有 UT/ATK 用例的 fused 均为 2 的幂，故未被覆盖。
tests/atk_tests/triton/mamba_mimo_bwd_fwd/init.py — no issues（仅 docstring）。
tests/atk_tests/triton/mamba_mimo_bwd_fwd/generate_mamba_mimo_bwd_fwd.py — no issues：17 个 shape 与 yaml 输入顺序一一对应，注册名与 yaml generate 字段一致，_case_index 类属性循环模式与仓库内既有 generator 一致。
tests/atk_tests/triton/mamba_mimo_bwd_fwd/mamba_mimo_bwd_fwd.yaml — no issues：triton_name/api_type/triton_api_type/generate 均与代码注册名匹配，single_bm/bf16_error/bf16_eb 键结构与既有 ATK 配置一致，无键拼写错误。
tests/atk_tests/triton/mamba_mimo_bwd_fwd/triton_mamba_mimo_bwd_fwd.py — no issues：_arguments 的 19 个实参顺序与 golden/kernel 签名一致，import 风格与既有 ATK 用例一致。
tests/unit_tests/triton/mamba_mimo_bwd_fwd_reference.py — no issues：逐项核对（factor/gamma、rotary、diagonal、inter/intra、state 递推、dz/dmimo_z/dmimo_o 约简、permute 维度）均与 kernel 语义一致。
tests/unit_tests/triton/test_mamba_mimo_bwd_fwd.py — no issues：4 种 D/Z/reduceO 组合、GQA、尾 chunk、5 组 shape 参数、非法入参用例均合理。

总体风险评估

发现数量：P0 0 个、P1 0 个、P2 1 个、P3 0 个。
总体判断：该 PR 主体实现与 CPU FP32 golden 的数学语义高度一致，测试覆盖较完整（两条 kernel 路径、GQA、尾 chunk、可选输入组合均有覆盖），ATK 配置与既有仓库约定一致。唯一确认的问题是一个仅在 chunk_size*rank 非 2 的幂时触发的越界写并发竞争（P2），属于边界配置下的潜在数据损坏，不影响当前已测试的全部 2 的幂配置；建议修复并补充非 2 的幂用例后再合入。整体风险可控。

类型	数量
🔴 阻塞	0
🟡 建议	1

💬 仅评论

不准确？

ascend-robot成员

8月19日添加了label：ascend-cla/yes

ascend-robot成员

8月19日评论：

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

⚠️ This PR does not yet meet the following requirements:lgtm (requires ≥ 2 person(s) per module)、approve (requires ≥ 1 person(s) per module)

Module Approval Details

module	lgtm status	approve status
docs	❌ (0/2)(You can also ask: 孙银磊, gitcode_lhp, zjchenn, 华郁秀, 周蓓蓉)	❌ (0/1)(You can also ask: 刘哲续, 孙银磊, bigdog1206, 郑加利, 刘荣)
repo-Ascend/MindSpeed-Ops	❌ (0/2)(You can also ask: 华郁秀, cjy840282, 周蓓蓉, guihaowen666, HANHU1CHEN)	❌ (0/1)(You can also ask: 朱彦儒, 雷珍珍, bigdog1206, 周蓓蓉, 刘荣)

💡 Tip:

Committer can comment /approve or /lgtm

Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

hahahascend, thanks for your pull request. All authors of the commits have signed the CLA. 👍

atomgit-bot8月19日进行代码检视2

tests/unit_tests/triton/mamba_mimo_bwd_fwd_reference.py

133+ dmimo_o = (output_gated * dout_h.unsqueeze(3)).sum(dim=2)

134+ d_pre_gate = dout_h.unsqueeze(3) * mimo_o.float()[None, :, None, :, :] * output if has_z else None

135+ else:

136+ dout_h = _pad(dout, padded_seqlen, 1).float().permute(0, 3, 1, 2, 4)

atomgit-bot8月19日评论：

🟡 Medium Priority

变更行：mamba_mimo_bwd_fwd_reference.py 第 136 行 dout_h = _pad(dout, padded_seqlen, 1).float().permute(0, 3, 1, 2, 4)。

证据链：

建议：将 reduce_o=False 分支的 permute 改为与 reduce_o=True 分支一致的 permute(0, 2, 1, 3)，并补充一个 mimo_o=None 的 UT 用例验证该分支（若 PR 声称已覆盖 reduceO 组合，需确认该用例实际执行通过）。

改动建议

137	- dout_h = _pad(dout, padded_seqlen, 1).float().permute(0, 3, 1, ~~2, 4~~)
137	+ dout_h = _pad(dout, padded_seqlen, 1).float().permute(0, 2, 1, 3)

应用建议

不准确？

hahahascend

9月1日评论：

已修改

H hahahascend

8月19日推送 1 个提交:910e7748-fix: handle padded Mamba MIMO output rows

ascend-robot成员

8月19日删除了label：ascend-cla/yes

ascend-robot成员

8月19日添加了label：stat/needs-squash

ascend-robot成员

8月19日添加了label：ascend-cla/yes

LinShua成员9月15日进行代码检视1

.pre-commit-config.yaml

4040 args: [

4141 "-L",

42- "CANN,cann,NNAL,nnal,ASCEND,ascend,EnQue,CopyIn,ArchType,AND,ND,tbe,copyin,alog",

42+ "CANN,cann,NNAL,nnal,ASCEND,ascend,EnQue,CopyIn,ArchType,AND,ND,tbe,copyin,alog,DOUT,dout",

LinShua9月15日评论：

需要补充下以下信息：
（1）ATK的性能和精度通过截图
（2）GPU上pytorch小算子与GPU上开源triton算子的精度对齐信息（以确保pytorch小算子标杆的准确性）