Star15
90
代码
代码
Issues34
Pull Requests41
Actions
项目讨论
Wiki
分析
Star15
90
已开启
feat: add Triton-Ascend implementation of Mamba-3 mamba_mimo_bwd_fwd #125
hahahascend创建于 8月19日
feat: add Triton-Ascend implementation of Mamba-3 mamba_mimo_bwd_fwd #125
已开启
hahahascend创建于 8月19日
从hahahascend/MindSpeed-Opscodex/mamba-mimo-bwd-fwd-triton合入到Ascend/MindSpeed-Opsmaster
共 14 个文件变更+2085-1
M.pre-commit-config.yaml+1-1
@@ -39,7 +39,7 @@ repos:
39 - id: codespell39 - id: codespell
40 args: [40 args: [
41 "-L",41 "-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",
L
LLinShua9月15日

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

likedislike
43 "--skip",43 "--skip",
44 "*.toml,*.py,*.cpp,*.hpp,*.c,*.h",44 "*.toml,*.py,*.cpp,*.hpp,*.c,*.h",
45 ]45 ]
MREADME.md+1-0
@@ -100,6 +100,7 @@ python -c "from mindspeed_ops.api.tilelang.sinkhorn import sinkhorn; print(sinkh
100| [`FusedCrossEntropyLoss`](docs/triton/fused_cross_entropy_loss.md) | 自定义融合算子 | ❌ | √ | ❌ |100| [`FusedCrossEntropyLoss`](docs/triton/fused_cross_entropy_loss.md) | 自定义融合算子 | ❌ | √ | ❌ |
101| [`RmsnormWithoutWeight`](docs/triton/rmsnorm_without_weight.md) | 自定义融合算子 | ❌ | √ | ❌ |101| [`RmsnormWithoutWeight`](docs/triton/rmsnorm_without_weight.md) | 自定义融合算子 | ❌ | √ | ❌ |
102| [`Sinkhorn`](docs/triton/sinkhorn.md) | Sinkhorn算法 | ❌ | √ | √ |102| [`Sinkhorn`](docs/triton/sinkhorn.md) | Sinkhorn算法 | ❌ | √ | √ |
103+| [`mamba_mimo_bwd_fwd`](docs/triton/mamba_mimo_bwd_fwd.md) | Mamba-3 MIMO backward 第一阶段 | ❌ | √ | ❌ |
103| [`Causal_conv1d`](docs/triton/causal_conv1d.md) | causal_conv1d算子 | ❌ | √ | ❌ |104| [`Causal_conv1d`](docs/triton/causal_conv1d.md) | causal_conv1d算子 | ❌ | √ | ❌ |
104| [`Wy_fast`](docs/triton/wy_fast.md) | wy_fast算子 | ❌ | √ | ❌ |105| [`Wy_fast`](docs/triton/wy_fast.md) | wy_fast算子 | ❌ | √ | ❌ |
105| [`RmsNormGated`](docs/triton/rmsnormgated.md) | RmsNormGated算子 | ❌ | √ | ❌ |106| [`RmsNormGated`](docs/triton/rmsnormgated.md) | RmsNormGated算子 | ❌ | √ | ❌ |
Adocs/triton/mamba_mimo_bwd_fwd.md+126-0
@@ -0,0 +1,126 @@
1+# mamba_mimo_bwd_fwd Triton 算子
2+ 
3+## 算子概述
4+ 
5+`mamba_mimo_bwd_fwd` 是 Mamba-3 MIMO backward 的第一阶段。实现以
6+state-spaces/mamba 提交 `48582f0592f1d55b441cd963aa5f9226b083f02a` 为语义基线，将其
7+TileLang 算法使用 Triton-Ascend 重写，用于 Ascend arch32 NPU。
8+ 
9+算子重新执行 chunk scan，累加 MIMO 输出投影梯度，并生成后续 backward 使用的 recurrent state
10+和同 token QK 点积。支持 D/Z 可选分支、两种 DOUT 布局、GQA 和尾 chunk。
11+ 
12+## 使用方法
13+ 
14+```python
15+from mindspeed_ops.api.triton import mamba_mimo_bwd_fwd
16+ 
17+outputs = mamba_mimo_bwd_fwd(
18+ dout,
19+ q,
20+ k,
21+ v,
22+ q_bias,
23+ k_bias,
24+ mimo_v,
25+ mimo_o,
26+ z,
27+ mimo_z,
28+ angles,
29+ da_cs,
30+ da_cs_rev,
31+ dt,
32+ trap,
33+ d,
34+ segsum,
35+ chunk_size=16,
36+ rotary_dim_divisor=4,
37+)
38+```
39+ 
40+`mimo_o is not None` 选择 reduced DOUT；`z is not None` 启用 `DZ/DMIMO_Z`；
41+`d is not None` 启用 skip 项。
42+ 
43+## 输入与输出
44+ 
45+符号：batch `B`、序列长度 `S`、value head 数 `H`、Q/K group 数 `G`、MIMO rank `R`、
46+state dimension `N`、value dimension `P`、chunk size `C`。
47+ 
48+| 输入 | Shape | dtype |
49+| --- | --- | --- |
50+| `dout` | reduced: `[B,S,H,P]`; non-reduced: `[B,S,R,H,P]` | BF16 |
51+| `q`, `k` | `[B,S,R,G,N]` | BF16 |
52+| `v` | `[B,S,H,P]` | BF16 |
53+| `q_bias`, `k_bias` | `[H,R,N]` | FP32 |
54+| `mimo_v`, optional `mimo_o`, optional `mimo_z` | `[H,R,P]` | FP32 |
55+| optional `z` | `[B,S,H,P]` | BF16 |
56+| `angles` | `[B,S,H,N/rotary_dim_divisor]` | FP32 |
57+| `da_cs`, `da_cs_rev`, `dt` | `[B,H,S]` | FP32 |
58+| `trap` | `[B,H,S]` | BF16 |
59+| optional `d` | `[H]` | FP32 |
60+| `segsum` | `[B,H,ceil(S/C),C,C]` | FP32 |
61+ 
62+| 输出 | Shape | dtype |
63+| --- | --- | --- |
64+| `dmimo_o` | `[B,H,R,P]` 或 `None` | FP32 |
65+| `states` | `[B,H,ceil(S/C),N,P]` | BF16 |
66+| `dz` | `[B,S,H,P]` 或 `None` | BF16 |
67+| `dmimo_z` | `[B,H,R,P]` 或 `None` | FP32 |
68+| `qk_dot` | `[B,H,S,R,R]` | BF16 |
69+ 
70+约束：`H % G == 0`、`N/P <= 128`、`C*R <= 128`，`N` 为偶数且可被 rotary divisor
71+整除。内部 recurrent state、dot 和梯度累加使用 FP32，写出 dtype 与上游接口一致。
72+ 
73+## Triton 实现
74+ 
75+实现由以下 Triton kernel 组成：
76+ 
77+1. token、rotary、causal 三个全序列 control kernel：用连续大块代替逐 chunk 的稀疏 lane；
78+2. `mamba_mimo_qk_bias_kernel`：同一 program 同时生成带 BF16 舍入的 biased Q/K；
79+3. bias kernel 内先复制/缩放全维，再由 `mamba_mimo_qk_rotary_patch_kernel`
80+ 只覆盖两个 rotary 列区间，避免在非 rotary 列执行复杂路径；
81+4. `mamba_mimo_psi_kernel`：生成 `PsiV`；
82+5. `mamba_mimo_state_scan_kernel`：每个 batch-head 持久化 FP32 recurrent state；
83+6. `mamba_mimo_output_kernel`：所有 chunk 并行计算 BF16 dot（FP32 累加）和输出，value 维按最多 64 列分块；
84+7. gather、projection-grad 和 DZ kernel：提取同 token QK block 并完成可选梯度分支。
85+ 
86+`N>64` 或 `R=1` 时自动使用分块 preprocess fallback，以规避 192 KiB UB 上限和当前
87+Triton-Ascend 的 unit-stride 编译缺陷。主输入和中间 fragment 按上游 BF16 边界舍入，dot 使用 FP32 累加，
88+recurrent state 和投影梯度保持 FP32 累积。对角项和 DZ 按 rank 顺序执行 BF16 舍入，
89+避免将数学等价但数值不同的 FP32 路径作为上游标杆。
90+ 
91+## 精度验证
92+ 
93+UT/ATK 使用显式复现上游 BF16 舍入的 PyTorch reference，原 FP32 reference 保留用于诊断。
94+UT 指标为 `RMS(actual - reference) / RMS(reference) < 0.001`，阈值未放宽。
95+ 
96+2026-10-17，Ascend910_9382 / CANN 9.0：**20 个 UT 通过**。
97+同一输入文件下，20 组配置的 76 项有效输出均通过；NPU 对 A40 原始上游 TileLang
98+kernel 的最大 relative RMS 为 `0.000234260140`（约 0.023426%）。
99+A40 CUDA BF16 reference 对原始 GPU kernel 三次重复的 228 项比较也全部通过。
100+ 
101+```bash
102+pytest -q tests/unit_tests/triton/test_mamba_mimo_bwd_fwd.py
103+```
104+ 
105+## NPU 性能记录
106+ 
107+当前 BF16 版本形状为 `B=1,S=8192,H=16,G=1,N=P=64,R=4,C=16`，
108+D/Z/reduceO 全开。三个独立进程各预热 10 次、测量 30 次同步 API wall time，
109+包含输出/workspace 分配，排除编译、输入生成、H2D 和第二阶段。
110+ 
111+CANN 9.0 的进程中位数为 53.817860 / 53.786125 / 53.915635 ms，
112+全部 90 个样本中位数为 **53.849475 ms**，样本标准差 0.127761 ms。
113+ 
114+历史 H100 SXM 数据为 8.201160 ms，与当前耗时之比为 0.152298。两者形状和计时边界
115+一致，但尚无历史 H100 输入文件用于核实相同输入。当前 BF16 dot 改变了计算路径，
116+需要重新确认 bound 分类，故此比值仅供参考，不能直接宣布 NVIDIA 折算性能验收通过。
117+旧版 CANN 8.5 的 59.592497 ms 属于历史 FP32 中间计算版本，不代表当前实现。
118+ 
119+## ATK
120+ 
121+配置位于 `tests/atk_tests/triton/mamba_mimo_bwd_fwd/`，原始精度标准未修改。
122+ATK 26.7.8 实测：精度 4/4 Pass；性能 4/4 Pass，平均性能比 2.2977。
123+性能比较对象为**同机 PyTorch NPU BF16 reference**，不代表 H100 性能验收。
124+ 
125+完整环境、逐用例数据、原始报告路径和数值语义说明见
126+[本轮评审验证记录](mamba_mimo_bwd_fwd_review_validation.md)。
Adocs/triton/mamba_mimo_bwd_fwd_review_validation.md+83-0
@@ -0,0 +1,83 @@
1+# PR #125：BF16 标杆与 NPU 复验
2+ 
3+## 数值语义变更
4+ 
5+原 PR 使用 FP32 中间计算的数学 reference。本次另加
6+`tests/unit_tests/triton/mamba_mimo_bwd_fwd_bf16_reference.py`，显式复现上游
7+`48582f0592f1d55b441cd963aa5f9226b083f02a` TileLang 第一阶段的 BF16 舍入边界。
8+原 FP32 reference 保留用于数学诊断；UT/ATK 使用新的混合精度 reference。
9+这不是放宽阈值：相对 RMS 阈值仍为严格 `<0.001`。
10+ 
11+上游 GPU 实现是 TileLang，NPU 实现是 Triton-Ascend，报告中应明确区分。
12+ 
13+## A40 已核验证据
14+ 
15+来自 `mamba_pr125_bf16_reference.tar.gz` 的 20 组配置使用原 PR 输入公式和
16+seed=123/0/1/42。独立重算三次共 228 项有效输出误差均通过，最大值
17+`0.00023416845025872978`，对应 `tail_gqa_full_seed0/DZ`。
18+上游 kernel 与固定 commit 逐字节一致；747 个文件哈希全部匹配；20 组输入
19+与前轮失败报告完全相同。每组的三次输出稳定，没有通过替换输入或修改 kernel 获得通过。
20+ 
21+本机还执行了 CPU BF16 reference，与保存的 CUDA reference 比较，76 项全部通过，
22+最大 relative RMS `0.00004576264458687435`。
23+本机没有重新执行 CUDA kernel，GPU 结果的核验依赖包内源码、原始 tensor 和日志。
24+ 
25+## NPU 实现修改
26+ 
27+- 在 bias、旋转、PsiV、trap scaling、state 输入、causal QK、output 和 gating
28+ 各阶段补齐上游 BF16 舍入；矩阵乘法和持续 state 仍使用 FP32 累积。
29+- 对角项按 rank 顺序进行 BF16 乘法和加法舍入，不能合并为 FP32 dot。
30+- output kernel 沿 value 维分块，最多 64 列，避免 `P=128` 时 UB 超限。
31+- ATK golden 入口明确恢复契约 dtype，防止工具自动升精改变 BF16 reference 的定义。
32+- 增加尾块 seed 和多 batch/group 回归用例。
33+ 
34+## 本机复验环境与证据位置
35+ 
36+Ascend910_9382、CANN 9.0、PyTorch 2.9.0+cpu 配合 torch_npu、Triton 3.2.0。
37+本机证据目录：`/host-model/slai/user-wangyihan/dzl/mamba-pr125-evidence/`。
38+该目录保存 ATK 工具、运行脚本、原始日志和 Excel；不作为仓库运行时依赖。
39+ 
40+## 最终复验结果（2026-10-17）
41+ 
42+| 检查 | 结果 |
43+| --- | --- |
44+| UT | 20 passed |
45+| NPU vs A40 CUDA BF16 reference，同一输入文件 | 76/76 输出通过；最大 relative RMS 0.000113489044 |
46+| NPU vs A40 原始上游 TileLang kernel，同一输入文件 | 76/76 输出通过；最大 relative RMS 0.000234260140 |
47+| ATK 精度，CPU BF16 reference | 4/4，通过率 100%，Pass |
48+| ATK 性能，同机 PyTorch NPU BF16 reference | 4/4，平均性能比 2.2977，Pass |
49+ 
50+ATK 原始 YAML 和精度标准未修改。性能逐用例数据如下，单位为微秒：
51+ 
52+| ID | Triton NPU | PyTorch NPU | PyTorch / Triton |
53+| --- | ---: | ---: | ---: |
54+| 0 | 388.658 | 601.002 | 1.54635 |
55+| 1 | 398.752 | 930.842 | 2.33439 |
56+| 2 | 400.264 | 860.775 | 2.15052 |
57+| 3 | 410.348 | 1296.523 | 3.15957 |
58+ 
59+原始日志为 `bf16_ut.log`、`bf16_a40_npu.log`、`bf16_atk_accuracy.log` 和
60+`bf16_atk_performance.log`；逐输出数据为 `compare_a40_npu.json`，输入和输出为
61+`a40_npu_cases/`。ATK Excel 位于证据目录的以下子目录：
62+ 
63+- `atk_output/all_mamba_mimo_bwd_fwd_2026-10-17-17-12-29-855283/report/`（精度）
64+- `atk_output/all_mamba_mimo_bwd_fwd_2026-10-17-17-13-38-223099/report/`（性能）
65+ 
66+`screenshots/` 是浏览器渲染上述 Excel 数据的截图，表头翻译为英文，测量值不变；
67+不是 ATK 原生界面截图。原始 Excel 随包提供，截图注明源文件及 SHA256。
68+ 
69+## 大形状性能及边界
70+ 
71+`B=1,S=8192,H=16,G=1,N=P=64,R=4,C=16`，启用 D/Z/reduceO。
72+每进程编译和 10 次预热后，采集 30 次同步 API wall time，包含分配，排除输入生成、
73+H2D 和第二阶段。三个独立进程中位数为 53.817860、53.786125、53.915635 ms；
74+90 个样本的总中位数为 **53.849475 ms**，范围 53.730770–54.718640 ms，
75+样本标准差 0.127761 ms。原始数据为 `bf16_performance_0/1/2.json`。
76+ 
77+历史 H100 中位数 8.201160 ms 除以当前耗时为 0.152298。H100 历史输入文件尚未取得，
78+只能注明形状和计时边界一致，不能声称输入逐位相同。当前引入 BF16 dot，需重新确认
79+实际 bound 分类及硬件规格映射；不能直接沿用旧版 CV Vector-bound 分类宣布性能验收。
80+ATK 性能 Pass 的比较对象是同机 PyTorch NPU，不能代替 NVIDIA 折算验收。
81+ 
82+本轮代码变更基于 PR 提交 `910e77489bef33d87de02d432f28c27bfacecd86`。
83+原始测量和报告在本轮修改后的工作树上生成。
Mmindspeed_ops/api/triton/__init__.py+3-0
@@ -0,0 +1,3 @@
1+from mindspeed_ops.api.triton.mamba_mimo_bwd_fwd import MambaMimoBwdFwdOutputs, mamba_mimo_bwd_fwd
2+ 
3+__all__ = ["MambaMimoBwdFwdOutputs", "mamba_mimo_bwd_fwd"]
Amindspeed_ops/api/triton/mamba_mimo_bwd_fwd.py+211-0
@@ -0,0 +1,211 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+# pylint: disable=import-outside-toplevel,too-many-arguments,too-many-locals
3+ 
4+"""Public Triton-Ascend API for Mamba-3 ``mamba_mimo_bwd_fwd``."""
5+ 
6+from __future__ import annotations
7+from functools import lru_cache
8+ 
9+from typing import NamedTuple
10+ 
11+import torch
12+ 
13+ 
14+__all__ = ["MambaMimoBwdFwdOutputs", "mamba_mimo_bwd_fwd"]
15+ 
16+ 
17+class MambaMimoBwdFwdOutputs(NamedTuple):
18+ """Accuracy-critical outputs of the first Mamba MIMO backward phase."""
19+ 
20+ dmimo_o: torch.Tensor | None
21+ states: torch.Tensor
22+ dz: torch.Tensor | None
23+ dmimo_z: torch.Tensor | None
24+ qk_dot: torch.Tensor
25+ 
26+ 
27+@lru_cache
28+def _is_arch35_device(device: torch.device) -> bool:
29+ """Check the architecture of the device that owns the input tensor."""
30+ import torch_npu
31+ 
32+ device_name = torch_npu.npu.get_device_properties(device).name
33+ return "Ascend910_95" in device_name or "Ascend950" in device_name
34+ 
35+ 
36+def _check_tensor(
37+ name: str,
38+ tensor: torch.Tensor,
39+ shape: tuple[int, ...],
40+ dtype: torch.dtype,
41+ device: torch.device,
42+) -> None:
43+ if tensor.shape != shape:
44+ raise ValueError(f"{name} must have shape {list(shape)}, got {list(tensor.shape)}")
45+ if tensor.dtype != dtype:
46+ raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}")
47+ if tensor.device != device:
48+ raise ValueError(f"{name} must be on {device}, got {tensor.device}")
49+ 
50+ 
51+def _validate_inputs(
52+ dout: torch.Tensor,
53+ q: torch.Tensor,
54+ k: torch.Tensor,
55+ v: torch.Tensor,
56+ q_bias: torch.Tensor,
57+ k_bias: torch.Tensor,
58+ mimo_v: torch.Tensor,
59+ mimo_o: torch.Tensor | None,
60+ z: torch.Tensor | None,
61+ mimo_z: torch.Tensor | None,
62+ angles: torch.Tensor,
63+ da_cs: torch.Tensor,
64+ da_cs_rev: torch.Tensor,
65+ dt: torch.Tensor,
66+ trap: torch.Tensor,
67+ d: torch.Tensor | None,
68+ segsum: torch.Tensor,
69+ chunk_size: int,
70+ rotary_dim_divisor: int,
71+) -> None:
72+ if q.dim() != 5:
73+ raise ValueError("q must have shape [B, S, R, G, N]")
74+ batch, seqlen, rank, groups, state_dim = q.shape
75+ if batch < 1 or seqlen < 1 or rank < 1 or groups < 1:
76+ raise ValueError("B, S, R and G must be positive")
77+ if q.dtype != torch.bfloat16:
78+ raise TypeError("q, k, v, dout and trap must use torch.bfloat16")
79+ if q.device.type not in ("npu", "privateuseone"):
80+ raise ValueError("mamba_mimo_bwd_fwd requires NPU tensors")
81+ if _is_arch35_device(q.device):
82+ raise NotImplementedError("mamba_mimo_bwd_fwd Triton currently supports arch32 only")
83+ if chunk_size < 1 or chunk_size * q.shape[2] > 128:
84+ raise ValueError("chunk_size * R must be in [1, 128]")
85+ if rotary_dim_divisor not in (2, 4):
86+ raise ValueError("rotary_dim_divisor must be 2 or 4")
87+ 
88+ if state_dim % rotary_dim_divisor or state_dim % 2:
89+ raise ValueError("N must be even and divisible by rotary_dim_divisor")
90+ if state_dim > 128:
91+ raise ValueError("N greater than 128 is not supported")
92+ if k.shape != q.shape:
93+ raise ValueError("k must have the same shape as q")
94+ if v.dim() != 4 or v.shape[:2] != (batch, seqlen):
95+ raise ValueError("v must have shape [B, S, H, P]")
96+ heads, value_dim = v.shape[2:]
97+ if heads % groups:
98+ raise ValueError("H must be divisible by G")
99+ if value_dim > 128:
100+ raise ValueError("P greater than 128 is not supported")
101+ 
102+ device = q.device
103+ _check_tensor("k", k, tuple(q.shape), torch.bfloat16, device)
104+ _check_tensor("v", v, (batch, seqlen, heads, value_dim), torch.bfloat16, device)
105+ _check_tensor("q_bias", q_bias, (heads, rank, state_dim), torch.float32, device)
106+ _check_tensor("k_bias", k_bias, (heads, rank, state_dim), torch.float32, device)
107+ _check_tensor("mimo_v", mimo_v, (heads, rank, value_dim), torch.float32, device)
108+ _check_tensor(
109+ "angles",
110+ angles,
111+ (batch, seqlen, heads, state_dim // rotary_dim_divisor),
112+ torch.float32,
113+ device,
114+ )
115+ for name, tensor in (("da_cs", da_cs), ("da_cs_rev", da_cs_rev), ("dt", dt)):
116+ _check_tensor(name, tensor, (batch, heads, seqlen), torch.float32, device)
117+ _check_tensor("trap", trap, (batch, heads, seqlen), torch.bfloat16, device)
118+ nchunks = (seqlen + chunk_size - 1) // chunk_size
119+ _check_tensor("segsum", segsum, (batch, heads, nchunks, chunk_size, chunk_size), torch.float32, device)
120+ if d is not None:
121+ _check_tensor("d", d, (heads,), torch.float32, device)
122+ 
123+ if mimo_o is not None:
124+ _check_tensor("mimo_o", mimo_o, (heads, rank, value_dim), torch.float32, device)
125+ _check_tensor("dout", dout, (batch, seqlen, heads, value_dim), torch.bfloat16, device)
126+ else:
127+ _check_tensor("dout", dout, (batch, seqlen, rank, heads, value_dim), torch.bfloat16, device)
128+ 
129+ if z is None:
130+ if mimo_z is not None:
131+ raise ValueError("mimo_z must be None when z is None")
132+ else:
133+ _check_tensor("z", z, (batch, seqlen, heads, value_dim), torch.bfloat16, device)
134+ if mimo_z is None:
135+ raise ValueError("mimo_z is required when z is present")
136+ _check_tensor("mimo_z", mimo_z, (heads, rank, value_dim), torch.float32, device)
137+ 
138+ 
139+def mamba_mimo_bwd_fwd(
140+ dout: torch.Tensor,
141+ q: torch.Tensor,
142+ k: torch.Tensor,
143+ v: torch.Tensor,
144+ q_bias: torch.Tensor,
145+ k_bias: torch.Tensor,
146+ mimo_v: torch.Tensor,
147+ mimo_o: torch.Tensor | None,
148+ z: torch.Tensor | None,
149+ mimo_z: torch.Tensor | None,
150+ angles: torch.Tensor,
151+ da_cs: torch.Tensor,
152+ da_cs_rev: torch.Tensor,
153+ dt: torch.Tensor,
154+ trap: torch.Tensor,
155+ d: torch.Tensor | None,
156+ segsum: torch.Tensor,
157+ chunk_size: int = 16,
158+ rotary_dim_divisor: int = 4,
159+) -> MambaMimoBwdFwdOutputs:
160+ """Run the issue-22 Mamba-3 MIMO backward-forward phase on arch32.
161+ 
162+ The implementation follows state-spaces/mamba commit ``48582f0``. Passing
163+ ``mimo_o`` selects reduced DOUT; passing ``z`` enables ``dz`` and
164+ ``dmimo_z``. BF16 fragment roundings follow the upstream GPU operator;
165+ matrix products and recurrent-state accumulation use FP32 accumulation.
166+ """
167+ _validate_inputs(
168+ dout,
169+ q,
170+ k,
171+ v,
172+ q_bias,
173+ k_bias,
174+ mimo_v,
175+ mimo_o,
176+ z,
177+ mimo_z,
178+ angles,
179+ da_cs,
180+ da_cs_rev,
181+ dt,
182+ trap,
183+ d,
184+ segsum,
185+ chunk_size,
186+ rotary_dim_divisor,
187+ )
188+ from mindspeed_ops.arch32.triton.mamba_mimo_bwd_fwd import mamba_mimo_bwd_fwd_triton
189+ 
190+ outputs = mamba_mimo_bwd_fwd_triton(
191+ dout.contiguous(),
192+ q.contiguous(),
193+ k.contiguous(),
194+ v.contiguous(),
195+ q_bias.contiguous(),
196+ k_bias.contiguous(),
197+ mimo_v.contiguous(),
198+ None if mimo_o is None else mimo_o.contiguous(),
199+ None if z is None else z.contiguous(),
200+ None if mimo_z is None else mimo_z.contiguous(),
201+ angles.contiguous(),
202+ da_cs.contiguous(),
203+ da_cs_rev.contiguous(),
204+ dt.contiguous(),
205+ trap.contiguous(),
206+ None if d is None else d.contiguous(),
207+ segsum.contiguous(),
208+ chunk_size,
209+ rotary_dim_divisor,
210+ )
211+ return MambaMimoBwdFwdOutputs(*outputs)
Atests/atk_tests/triton/mamba_mimo_bwd_fwd/__init__.py+1-0
@@ -0,0 +1 @@
1+"""ATK Triton coverage for mamba_mimo_bwd_fwd."""
Atests/atk_tests/triton/mamba_mimo_bwd_fwd/generate_mamba_mimo_bwd_fwd.py+45-0
@@ -0,0 +1,45 @@
1+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
2+ 
3+from atk.case_generator.generator.base_generator import CaseGenerator
4+from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
5+from atk.configs.case_config import CaseConfig
6+ 
7+ 
8+CASES = [
9+ (1, 16, 1, 1, 64, 64, 4, 16),
10+ (1, 32, 1, 1, 64, 64, 4, 16),
11+ (1, 20, 1, 1, 64, 64, 4, 16),
12+ (1, 32, 2, 1, 64, 64, 4, 16),
13+]
14+ 
15+ 
16+@GENERATOR_REGISTRY.register("generate_mamba_mimo_bwd_fwd")
17+class MambaMimoBwdFwdGenerator(CaseGenerator):
18+ _case_index = 0
19+ 
20+ def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
21+ batch, seqlen, heads, groups, state_dim, value_dim, rank, chunk_size = CASES[self._case_index % len(CASES)]
22+ self.__class__._case_index += 1
23+ nchunks = (seqlen + chunk_size - 1) // chunk_size
24+ shapes = [
25+ [batch, seqlen, heads, value_dim],
26+ [batch, seqlen, rank, groups, state_dim],
27+ [batch, seqlen, rank, groups, state_dim],
28+ [batch, seqlen, heads, value_dim],
29+ [heads, rank, state_dim],
30+ [heads, rank, state_dim],
31+ [heads, rank, value_dim],
32+ [heads, rank, value_dim],
33+ [batch, seqlen, heads, value_dim],
34+ [heads, rank, value_dim],
35+ [batch, seqlen, heads, state_dim // 4],
36+ [batch, heads, seqlen],
37+ [batch, heads, seqlen],
38+ [batch, heads, seqlen],
39+ [batch, heads, seqlen],
40+ [heads],
41+ [batch, heads, nchunks, chunk_size, chunk_size],
42+ ]
43+ for input_config, shape in zip(case_config.inputs[:17], shapes):
44+ input_config.shape = shape
45+ return case_config
Atests/atk_tests/triton/mamba_mimo_bwd_fwd/mamba_mimo_bwd_fwd.yaml+41-0
@@ -0,0 +1,41 @@
1+api: pytorch
2+version: v2.1
3+name: torch_mamba_mimo_bwd_fwd
4+triton_name: triton_mamba_mimo_bwd_fwd.TritonMambaMimoBwdFwdApi
5+api_type: torch_mamba_mimo_bwd_fwd
6+triton_api_type: triton_mamba_mimo_bwd_fwd
7+generate: generate_mamba_mimo_bwd_fwd
8+dtype_numbers: 4
9+backward: false
10+standard:
11+ acc:
12+ single_bm:
13+ type: high_performance
14+ bf16_error: 0.001
15+ bf16_eb: 0.001
16+ perf: not_key
17+inputs:
18+ - { name: grad_output, type: tensor, required: true, dtypes: { values: [bf16] }, ranges: { valid: { values: [[-1, 1]] }, invalid: { values: [[-1, 1]] } }, shapes: { dim_numbers: { values: [4] }, max_length: 1048576 } }
19+ - { name: q, type: tensor, required: true, dtypes: { values: [bf16] }, ranges: { valid: { values: [[-1, 1]] }, invalid: { values: [[-1, 1]] } }, shapes: { dim_numbers: { values: [5] }, max_length: 1048576 } }
20+ - { name: k, type: tensor, required: true, dtypes: { values: [bf16] }, ranges: { valid: { values: [[-1, 1]] }, invalid: { values: [[-1, 1]] } }, shapes: { dim_numbers: { values: [5] }, max_length: 1048576 } }
21+ - { name: v, type: tensor, required: true, dtypes: { values: [bf16] }, ranges: { valid: { values: [[-1, 1]] }, invalid: { values: [[-1, 1]] } }, shapes: { dim_numbers: { values: [4] }, max_length: 1048576 } }
22+ - { name: q_bias, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
23+ - { name: k_bias, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
24+ - { name: mimo_v, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
25+ - { name: mimo_o, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
26+ - { name: z, type: tensor, required: true, dtypes: { values: [bf16] }, ranges: { valid: { values: [[-1, 1]] }, invalid: { values: [[-1, 1]] } }, shapes: { dim_numbers: { values: [4] }, max_length: 1048576 } }
27+ - { name: mimo_z, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
28+ - { name: angles, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[0, 1]] }, invalid: { values: [[0, 1]] } }, shapes: { dim_numbers: { values: [4] }, max_length: 1048576 } }
29+ - { name: da_cs, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
30+ - { name: da_cs_rev, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
31+ - { name: dt, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[0.01, 0.1]] }, invalid: { values: [[0.01, 0.1]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
32+ - { name: trap, type: tensor, required: true, dtypes: { values: [bf16] }, ranges: { valid: { values: [[-1, 1]] }, invalid: { values: [[-1, 1]] } }, shapes: { dim_numbers: { values: [3] }, max_length: 1048576 } }
33+ - { name: d, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [1] }, max_length: 1048576 } }
34+ - { name: segsum, type: tensor, required: true, dtypes: { values: [fp32] }, ranges: { valid: { values: [[-0.2, 0.2]] }, invalid: { values: [[-0.2, 0.2]] } }, shapes: { dim_numbers: { values: [5] }, max_length: 1048576 } }
35+ - name: chunk_size
36+ type: attr
37+ required: false
38+ dtypes: { values: [int] }
39+ ranges:
40+ valid: { values: [16] }
41+ invalid: { values: [16] }
Atests/atk_tests/triton/mamba_mimo_bwd_fwd/triton_mamba_mimo_bwd_fwd.py+62-0
@@ -0,0 +1,62 @@
1+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
2+# pylint: disable=duplicate-code
3+ 
4+import torch
5+ 
6+from atk.configs.dataset_config import InputDataset
7+from atk.tasks.api_execute import register
8+from atk.tasks.api_execute.base_api import BaseApi
9+from atk.tasks.api_execute.triton_base_api import TritonBaseApi
10+ 
11+ 
12+def _arguments(input_data: InputDataset):
13+ names = (
14+ "grad_output",
15+ "q",
16+ "k",
17+ "v",
18+ "q_bias",
19+ "k_bias",
20+ "mimo_v",
21+ "mimo_o",
22+ "z",
23+ "mimo_z",
24+ "angles",
25+ "da_cs",
26+ "da_cs_rev",
27+ "dt",
28+ "trap",
29+ "d",
30+ "segsum",
31+ )
32+ args = [input_data.kwargs[name] for name in names]
33+ args.append(input_data.kwargs.get("chunk_size", 16))
34+ args.append(4)
35+ return args
36+ 
37+ 
38+@register("torch_mamba_mimo_bwd_fwd")
39+class TorchMambaMimoBwdFwdApi(BaseApi):
40+ def __call__(self, input_data: InputDataset, with_output: bool = False):
41+ from tests.unit_tests.triton.mamba_mimo_bwd_fwd_bf16_reference import (
42+ mamba_mimo_bwd_fwd_bf16_reference as mamba_mimo_bwd_fwd_golden,
43+ )
44+ 
45+ args = _arguments(input_data)
46+ # ATK promotes golden inputs. This reference deliberately models the
47+ # original BF16 fragment boundaries, including parameter load dtypes.
48+ bf16_indices = {0, 1, 2, 3, 8, 14}
49+ for index, value in enumerate(args[:17]):
50+ if value is not None:
51+ args[index] = value.to(torch.bfloat16 if index in bf16_indices else torch.float32)
52+ return mamba_mimo_bwd_fwd_golden(*args)
53+ 
54+ 
55+@register("triton_mamba_mimo_bwd_fwd")
56+class TritonMambaMimoBwdFwdApi(TritonBaseApi):
57+ def __call__(self, input_data: InputDataset, with_output: bool = False):
58+ from mindspeed_ops.api.triton.mamba_mimo_bwd_fwd import (
59+ mamba_mimo_bwd_fwd,
60+ )
61+ 
62+ return mamba_mimo_bwd_fwd(*_arguments(input_data))
Atests/unit_tests/triton/mamba_mimo_bwd_fwd_bf16_reference.py+207-0
@@ -0,0 +1,207 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+"""PyTorch reference for the BF16 arithmetic of Mamba-3 MIMO backward stage 1.
3+ 
4+Target: state-spaces/mamba commit 48582f0592f1d55b441cd963aa5f9226b083f02a,
5+ops/tilelang/mamba3/mamba3_mimo_bwd.py:mamba_mimo_bwd_fwd.
6+ 
7+This is a mixed-precision numerical reference, not the PR's all-FP32 golden.
8+BF16 fragment/store boundaries are represented by explicit BF16 round trips;
9+FP32 storage between those boundaries keeps the arithmetic readable. Only torch
10+is required. Outputs remain on the input device, before any batch reduction.
11+Disable TF32 when checking against the upstream kernel on CUDA.
12+"""
13+ 
14+from __future__ import annotations
15+ 
16+import torch
17+from torch import Tensor
18+ 
19+Outputs = tuple[Tensor | None, Tensor, Tensor | None, Tensor | None, Tensor]
20+ 
21+ 
22+def _bf16(value: Tensor) -> Tensor:
23+ """Round to BF16, then expose that rounded value in FP32 arithmetic."""
24+ return value.to(torch.bfloat16).to(torch.float32)
25+ 
26+ 
27+def _pad(value: Tensor | None, length: int, dim: int) -> Tensor | None:
28+ if value is None or value.shape[dim] == length:
29+ return value
30+ shape = list(value.shape)
31+ shape[dim] = length - value.shape[dim]
32+ return torch.cat((value, value.new_zeros(shape)), dim=dim)
33+ 
34+ 
35+def _rotate(value: Tensor, angles: Tensor, divisor: int) -> Tensor:
36+ """Rotate paired regions in FP32 and store them as BF16 fragments."""
37+ width = value.shape[-1] // divisor
38+ half = value.shape[-1] // 2
39+ first, second = value[..., :width], value[..., half : half + width]
40+ cos, sin = angles.cos(), angles.sin()
41+ result = value.clone()
42+ result[..., :width] = _bf16(cos * first - sin * second)
43+ result[..., half : half + width] = _bf16(sin * first + cos * second)
44+ return result
45+ 
46+ 
47+def mamba_mimo_bwd_fwd_bf16_reference(
48+ dout: Tensor,
49+ q: Tensor,
50+ k: Tensor,
51+ v: Tensor,
52+ q_bias: Tensor,
53+ k_bias: Tensor,
54+ mimo_v: Tensor,
55+ mimo_o: Tensor | None,
56+ z: Tensor | None,
57+ mimo_z: Tensor | None,
58+ angles: Tensor,
59+ da_cs: Tensor,
60+ da_cs_rev: Tensor,
61+ dt: Tensor,
62+ trap: Tensor,
63+ d: Tensor | None,
64+ segsum: Tensor,
65+ chunk_size: int,
66+ rotary_dim_divisor: int = 4,
67+) -> Outputs:
68+ """Return DMIMO_O, STATES, DZ, DMIMO_Z, QK_DOT with upstream BF16 rounding.
69+ 
70+ Inputs and optional branches follow the PR golden signature. Q/K have shape
71+ [B,S,R,G,N], V/Z [B,S,H,P], and DOUT [B,S,H,P] when mimo_o is present or
72+ [B,S,R,H,P] otherwise. Main inputs, DOUT and TRAP are BF16; controls and
73+ projections are FP32. A missing Z requires missing MIMO_Z.
74+ 
75+ DMIMO_O/DMIMO_Z are FP32 [B,H,R,P]; STATES is BF16 [B,H,ceil(S/C),N,P];
76+ DZ is BF16 [B,S,H,P]; QK_DOT is BF16 [B,H,S,R,R]. Inactive outputs are None.
77+ No kernel, model import, file access, device transfer or batch reduction is
78+ performed here. Validation is documented in
79+ docs/triton/mamba_mimo_bwd_fwd_review_validation.md.
80+ """
81+ if q.ndim != 5 or v.ndim != 4:
82+ raise ValueError('q must be [B,S,R,G,N] and v must be [B,S,H,P]')
83+ batch, seqlen, rank, groups, state_dim = q.shape
84+ heads, value_dim = v.shape[-2:]
85+ if min(batch, seqlen, rank, groups, state_dim, heads, value_dim, chunk_size) <= 0:
86+ raise ValueError('dimensions and chunk_size must be positive')
87+ if heads % groups or rotary_dim_divisor not in (2, 4) or state_dim % rotary_dim_divisor:
88+ raise ValueError('heads must divide into groups; rotary divisor must be 2 or 4 and divide N')
89+ if (z is None) != (mimo_z is None):
90+ raise ValueError('z and mimo_z must either both be present or both be None')
91+ nchunks = (seqlen + chunk_size - 1) // chunk_size
92+ bf16_shapes = {
93+ 'q': (q, (batch, seqlen, rank, groups, state_dim)),
94+ 'k': (k, q.shape),
95+ 'v': (v, (batch, seqlen, heads, value_dim)),
96+ 'dout': (
97+ dout,
98+ (batch, seqlen, heads, value_dim) if mimo_o is not None else (batch, seqlen, rank, heads, value_dim),
99+ ),
100+ 'trap': (trap, (batch, heads, seqlen)),
101+ }
102+ fp32_shapes = {
103+ 'q_bias': (q_bias, (heads, rank, state_dim)),
104+ 'k_bias': (k_bias, (heads, rank, state_dim)),
105+ 'mimo_v': (mimo_v, (heads, rank, value_dim)),
106+ 'angles': (angles, (batch, seqlen, heads, state_dim // rotary_dim_divisor)),
107+ 'da_cs': (da_cs, (batch, heads, seqlen)),
108+ 'da_cs_rev': (da_cs_rev, (batch, heads, seqlen)),
109+ 'dt': (dt, (batch, heads, seqlen)),
110+ 'segsum': (segsum, (batch, heads, nchunks, chunk_size, chunk_size)),
111+ }
112+ if mimo_o is not None:
113+ fp32_shapes['mimo_o'] = (mimo_o, (heads, rank, value_dim))
114+ if z is not None:
115+ bf16_shapes['z'] = (z, v.shape)
116+ fp32_shapes['mimo_z'] = (mimo_z, (heads, rank, value_dim))
117+ if d is not None:
118+ fp32_shapes['d'] = (d, (heads,))
119+ for dtype, specs in ((torch.bfloat16, bf16_shapes), (torch.float32, fp32_shapes)):
120+ for name, (tensor, shape) in specs.items():
121+ if tensor.shape != shape:
122+ raise ValueError(f'{name}: expected shape {tuple(shape)}, got {tuple(tensor.shape)}')
123+ if tensor.dtype != dtype:
124+ raise TypeError(f'{name}: expected {dtype}, got {tensor.dtype}')
125+ if tensor.device != q.device:
126+ raise ValueError(f'{name}: expected device {q.device}, got {tensor.device}')
127+ 
128+ padded_seqlen = nchunks * chunk_size
129+ q, k, v, dout, z, angles = (_pad(tensor, padded_seqlen, 1) for tensor in (q, k, v, dout, z, angles))
130+ da_cs, da_cs_rev, dt, trap = (_pad(tensor, padded_seqlen, 2) for tensor in (da_cs, da_cs_rev, dt, trap))
131+ repeats = heads // groups
132+ qh = q.float().repeat_interleave(repeats, dim=3).permute(0, 3, 1, 2, 4)
133+ kh = k.float().repeat_interleave(repeats, dim=3).permute(0, 3, 1, 2, 4)
134+ # Upstream loads FP32 bias into BF16, then rounds the BF16 addition.
135+ qh = _bf16(qh + _bf16(q_bias)[None, :, None])
136+ kh = _bf16(kh + _bf16(k_bias)[None, :, None])
137+ qk_dot = _bf16(torch.einsum('bhtrn,bhtsn->bhtrs', qh, kh))
138+ angle_h = angles.permute(0, 2, 1, 3).unsqueeze(3)
139+ q_rot = _rotate(qh, angle_h, rotary_dim_divisor)
140+ k_rot = _rotate(kh, angle_h, rotary_dim_divisor)
141+ psi_v = _bf16(v.float().permute(0, 2, 1, 3).unsqueeze(3) * _bf16(mimo_v)[None, :, None])
142+ 
143+ # DT is loaded as BF16. Gamma is FP32; shifted gamma and trap scale are BF16.
144+ gamma = _bf16(dt) * trap.float().sigmoid()
145+ shifted_gamma = torch.zeros_like(gamma)
146+ shifted_gamma[..., :-1] = _bf16(_bf16(dt[..., 1:]) * torch.sigmoid(-trap.float()[..., 1:]))
147+ factor = _bf16(gamma + shifted_gamma)
148+ k_scaled = _bf16(k_rot * factor[..., None, None])
149+ state = torch.zeros((batch, heads, state_dim, value_dim), device=q.device, dtype=torch.float32)
150+ states, output_chunks = [], []
151+ steps = torch.arange(chunk_size, device=q.device).repeat_interleave(rank)
152+ strictly_causal = steps[:, None] > steps[None, :]
153+ 
154+ for chunk in range(nchunks):
155+ start, end = chunk * chunk_size, (chunk + 1) * chunk_size
156+ qc = q_rot[:, :, start:end].flatten(2, 3)
157+ kc = k_scaled[:, :, start:end].flatten(2, 3)
158+ pc = psi_v[:, :, start:end].flatten(2, 3)
159+ states.append(state.to(torch.bfloat16))
160+ inter = qc @ _bf16(state)
161+ inter = inter * da_cs[:, :, start:end].exp().repeat_interleave(rank, 2).unsqueeze(-1)
162+ causal_qk = qc @ kc.transpose(-1, -2)
163+ causal_qk = causal_qk * segsum[:, :, chunk].exp().repeat_interleave(rank, 2).repeat_interleave(rank, 3)
164+ causal_qk = _bf16(causal_qk.masked_fill(~strictly_causal, 0))
165+ intra = causal_qk @ pc
166+ diagonal = torch.zeros((batch, heads, chunk_size, rank, value_dim), device=q.device, dtype=torch.float32)
167+ for r in range(rank):
168+ # BF16 multiplication and BF16 addition round separately (not FP32 FMA).
169+ product = _bf16(qk_dot[:, :, start:end, :, r, None] * psi_v[:, :, start:end, r, :].unsqueeze(3))
170+ diagonal = _bf16(diagonal + product)
171+ diagonal = _bf16(diagonal * gamma[:, :, start:end, None, None])
172+ output = inter + intra + diagonal.flatten(2, 3)
173+ if d is not None:
174+ output = output + d[None, :, None, None] * pc
175+ # out_prereduced_shared is BF16 in both reduced and non-reduced Z branches.
176+ output_chunks.append(_bf16(output).reshape(batch, heads, chunk_size, rank, value_dim))
177+ k_state = _bf16(kc * da_cs_rev[:, :, start:end].exp().repeat_interleave(rank, 2).unsqueeze(-1))
178+ last_valid = min(end, seqlen) - 1
179+ state = state * da_cs[:, :, last_valid, None, None].exp() + k_state.transpose(-1, -2) @ pc
180+ 
181+ output = torch.cat(output_chunks, dim=2)
182+ if z is not None:
183+ z_h = z.float().permute(0, 2, 1, 3).unsqueeze(3)
184+ projected_z = z_h * mimo_z[None, :, None]
185+ if mimo_o is not None:
186+ dout_h = dout.float().permute(0, 2, 1, 3).unsqueeze(3)
187+ if z is not None:
188+ silu = projected_z * torch.sigmoid(projected_z)
189+ dmimo_o = (output * silu * dout_h).sum(dim=2)
190+ pre_gate_grad = _bf16(_bf16(dout_h * _bf16(mimo_o)[None, :, None]) * output)
191+ else:
192+ dmimo_o = (output * dout_h).sum(dim=2)
193+ else:
194+ dmimo_o = None
195+ if z is not None:
196+ pre_gate_grad = _bf16(dout.float().permute(0, 3, 1, 2, 4) * output)
197+ if z is not None:
198+ grad_z_projected = pre_gate_grad * torch.sigmoid(projected_z) * (1 + projected_z * torch.sigmoid(-projected_z))
199+ dz_h = torch.zeros((batch, heads, padded_seqlen, value_dim), device=q.device, dtype=torch.float32)
200+ for r in range(rank):
201+ # The product is FP32 here; only the accumulated DZ fragment is BF16.
202+ dz_h = _bf16(dz_h + grad_z_projected[:, :, :, r, :] * mimo_z[None, :, r, None, :])
203+ dmimo_z = (grad_z_projected * z_h).sum(dim=2)
204+ dz = dz_h[:, :, :seqlen].permute(0, 2, 1, 3).to(torch.bfloat16)
205+ else:
206+ dz = dmimo_z = None
207+ return (dmimo_o, torch.stack(states, dim=2), dz, dmimo_z, qk_dot[:, :, :seqlen].to(torch.bfloat16))
Atests/unit_tests/triton/mamba_mimo_bwd_fwd_reference.py+155-0
@@ -0,0 +1,155 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+# pylint: disable=duplicate-code,too-many-arguments,too-many-locals
3+ 
4+"""FP32 PyTorch reference for the Triton Mamba-3 MIMO backward phase."""
5+ 
6+from __future__ import annotations
7+ 
8+import torch
9+ 
10+ 
11+def _pad(tensor: torch.Tensor, target: int, dim: int) -> torch.Tensor:
12+ amount = target - tensor.shape[dim]
13+ if amount == 0:
14+ return tensor
15+ shape = list(tensor.shape)
16+ shape[dim] = amount
17+ return torch.cat((tensor, torch.zeros(shape, dtype=tensor.dtype, device=tensor.device)), dim=dim)
18+ 
19+ 
20+def _rotate(tensor: torch.Tensor, angles: torch.Tensor, divisor: int) -> torch.Tensor:
21+ rotary_dim = tensor.shape[-1] // divisor
22+ first = tensor[..., :rotary_dim]
23+ second_start = tensor.shape[-1] // 2
24+ second = tensor[..., second_start : second_start + rotary_dim]
25+ cos = torch.cos(angles).unsqueeze(3)
26+ sin = torch.sin(angles).unsqueeze(3)
27+ result = tensor.clone()
28+ result[..., :rotary_dim] = cos * first - sin * second
29+ result[..., second_start : second_start + rotary_dim] = sin * first + cos * second
30+ return result
31+ 
32+ 
33+def mamba_mimo_bwd_fwd_golden(
34+ dout: torch.Tensor,
35+ q: torch.Tensor,
36+ k: torch.Tensor,
37+ v: torch.Tensor,
38+ q_bias: torch.Tensor,
39+ k_bias: torch.Tensor,
40+ mimo_v: torch.Tensor,
41+ mimo_o: torch.Tensor | None,
42+ z: torch.Tensor | None,
43+ mimo_z: torch.Tensor | None,
44+ angles: torch.Tensor,
45+ da_cs: torch.Tensor,
46+ da_cs_rev: torch.Tensor,
47+ dt: torch.Tensor,
48+ trap: torch.Tensor,
49+ d: torch.Tensor | None,
50+ segsum: torch.Tensor,
51+ chunk_size: int,
52+ rotary_dim_divisor: int = 4,
53+):
54+ """Return outputs with FP32 math and upstream-compatible output dtypes."""
55+ batch, seqlen, rank, groups, state_dim = q.shape
56+ input_dtype = q.dtype
57+ heads, value_dim = v.shape[2:]
58+ nchunks = (seqlen + chunk_size - 1) // chunk_size
59+ padded_seqlen = nchunks * chunk_size
60+ repeats = heads // groups
61+ 
62+ q = _pad(q, padded_seqlen, 1).float()
63+ k = _pad(k, padded_seqlen, 1).float()
64+ v = _pad(v, padded_seqlen, 1).float()
65+ angles = _pad(angles, padded_seqlen, 1).float()
66+ da_cs = _pad(da_cs, padded_seqlen, 2).float()
67+ da_cs_rev = _pad(da_cs_rev, padded_seqlen, 2).float()
68+ dt = _pad(dt, padded_seqlen, 2).float()
69+ trap = _pad(trap, padded_seqlen, 2).float()
70+ 
71+ qh = q.repeat_interleave(repeats, dim=3).permute(0, 3, 1, 2, 4)
72+ kh = k.repeat_interleave(repeats, dim=3).permute(0, 3, 1, 2, 4)
73+ qh = qh + q_bias.float()[None, :, None, :, :]
74+ kh = kh + k_bias.float()[None, :, None, :, :]
75+ qk_dot = torch.einsum("bhtqn,bhtkn->bhtqk", qh, kh)
76+ angle_h = angles.permute(0, 2, 1, 3)
77+ q_rot = _rotate(qh, angle_h, rotary_dim_divisor)
78+ k_rot = _rotate(kh, angle_h, rotary_dim_divisor)
79+ 
80+ sigmoid_trap = torch.sigmoid(trap)
81+ gamma = dt * sigmoid_trap
82+ shifted = torch.zeros_like(gamma)
83+ shifted[..., :-1] = dt[..., 1:] * (1.0 - sigmoid_trap[..., 1:])
84+ factor = gamma + shifted
85+ k_scaled = k_rot * factor[:, :, :, None, None]
86+ psi_v = v.permute(0, 2, 1, 3).unsqueeze(3) * mimo_v.float()[None, :, None, :, :]
87+ 
88+ state = torch.zeros((batch, heads, state_dim, value_dim), dtype=torch.float32, device=q.device)
89+ states = []
90+ output_chunks = []
91+ causal = torch.tril(
92+ torch.ones((chunk_size, chunk_size), dtype=torch.bool, device=q.device),
93+ diagonal=-1,
94+ )
95+ for chunk in range(nchunks):
96+ start = chunk * chunk_size
97+ end = start + chunk_size
98+ states.append(state.clone())
99+ q_chunk = q_rot[:, :, start:end]
100+ k_chunk = k_scaled[:, :, start:end]
101+ psi_chunk = psi_v[:, :, start:end]
102+ inter = torch.einsum("bhlrn,bhnp->bhlrp", q_chunk, state)
103+ inter = inter * torch.exp(da_cs[:, :, start:end])[:, :, :, None, None]
104+ qk_causal = torch.einsum("bhlon,bhmrn->bhlomr", q_chunk, k_chunk)
105+ qk_causal = qk_causal * torch.exp(segsum[:, :, chunk])[:, :, :, None, :, None]
106+ qk_causal = qk_causal.masked_fill(~causal[None, None, :, None, :, None], 0.0)
107+ intra = torch.einsum("bhlomr,bhmrp->bhlop", qk_causal, psi_chunk)
108+ diagonal = torch.einsum("bhlor,bhlrp->bhlop", qk_dot[:, :, start:end], psi_chunk)
109+ diagonal = diagonal * gamma[:, :, start:end, None, None]
110+ output = inter + intra + diagonal
111+ if d is not None:
112+ output = output + d.float()[None, :, None, None, None] * psi_chunk
113+ output_chunks.append(output)
114+ k_state = k_chunk * torch.exp(da_cs_rev[:, :, start:end])[:, :, :, None, None]
115+ state = state * torch.exp(da_cs[:, :, end - 1])[:, :, None, None]
116+ state = state + torch.einsum("bhlrn,bhlrp->bhnp", k_state, psi_chunk)
117+ 
118+ output = torch.cat(output_chunks, dim=2)
119+ has_z = z is not None
120+ reduce_o = mimo_o is not None
121+ if has_z:
122+ z_h = _pad(z, padded_seqlen, 1).float().permute(0, 2, 1, 3)
123+ z_projected = z_h.unsqueeze(3) * mimo_z.float()[None, :, None, :, :]
124+ sigmoid_z = torch.sigmoid(z_projected)
125+ silu_z = z_projected * sigmoid_z
126+ silu_grad = sigmoid_z * (1.0 + z_projected * (1.0 - sigmoid_z))
127+ else:
128+ z_h = silu_z = silu_grad = None
129+ 
130+ if reduce_o:
131+ dout_h = _pad(dout, padded_seqlen, 1).float().permute(0, 2, 1, 3)
132+ output_gated = output * silu_z if has_z else output
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+ # Non-reduced DOUT is [B, S, R, H, P]; align it to [B, H, S, R, P].
137+ dout_h = _pad(dout, padded_seqlen, 1).float().permute(0, 3, 1, 2, 4)
atomgit-bot
atomgit-botatomgit-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)
应用建议
likedislike
不准确？
H
hahahascend
9月1日 评论：
138+ dmimo_o = None
139+ d_pre_gate = dout_h * output if has_z else None
140+ 
141+ if has_z:
142+ d_z_projected = d_pre_gate * silu_grad
143+ dz = (d_z_projected * mimo_z.float()[None, :, None, :, :]).sum(dim=3)
144+ dz = dz[:, :, :seqlen].permute(0, 2, 1, 3).to(input_dtype)
145+ dmimo_z = (d_z_projected * z_h.unsqueeze(3)).sum(dim=2)
146+ else:
147+ dz = dmimo_z = None
148+ 
149+ return (
150+ dmimo_o,
151+ torch.stack(states, dim=2).to(input_dtype),
152+ dz,
153+ dmimo_z,
154+ qk_dot[:, :, :seqlen].to(input_dtype),
155+ )
Atests/unit_tests/triton/test_mamba_mimo_bwd_fwd.py+172-0
@@ -0,0 +1,172 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+# pylint: disable=duplicate-code,no-value-for-parameter,too-many-locals
3+ 
4+import pytest
5+import torch
6+ 
7+from mindspeed_ops.api.triton.mamba_mimo_bwd_fwd import mamba_mimo_bwd_fwd
8+from mindspeed_ops.api.triton.utils import get_available_device
9+from tests.unit_tests.triton.mamba_mimo_bwd_fwd_bf16_reference import (
10+ mamba_mimo_bwd_fwd_bf16_reference as mamba_mimo_bwd_fwd_golden,
11+)
12+from tests.utils import get_err_ratio
13+ 
14+ 
15+OUTPUT_NAMES = ("DMIMO_O", "STATES", "DZ", "DMIMO_Z", "QK_DOT")
16+ 
17+ 
18+def _make_inputs(
19+ *,
20+ seed=123,
21+ batch=1,
22+ seqlen=32,
23+ heads=1,
24+ groups=1,
25+ state_dim=64,
26+ value_dim=64,
27+ rank=4,
28+ chunk_size=16,
29+ reduce_o=True,
30+ has_z=False,
31+ has_d=False,
32+):
33+ torch.manual_seed(seed)
34+ dtype = torch.bfloat16
35+ q = torch.randn((batch, seqlen, rank, groups, state_dim), dtype=dtype)
36+ k = torch.randn_like(q)
37+ v = torch.randn((batch, seqlen, heads, value_dim), dtype=dtype)
38+ dout_shape = (batch, seqlen, heads, value_dim) if reduce_o else (batch, seqlen, rank, heads, value_dim)
39+ dout = torch.randn(dout_shape, dtype=dtype)
40+ q_bias = torch.randn((heads, rank, state_dim), dtype=torch.float32) / 8
41+ k_bias = torch.randn_like(q_bias)
42+ mimo_v = torch.randn((heads, rank, value_dim), dtype=torch.float32) / 4
43+ mimo_o = torch.randn((heads, rank, value_dim), dtype=torch.float32) / 4 if reduce_o else None
44+ z = torch.randn((batch, seqlen, heads, value_dim), dtype=dtype) if has_z else None
45+ mimo_z = torch.randn((heads, rank, value_dim), dtype=torch.float32) / 4 if has_z else None
46+ angles = torch.rand((batch, seqlen, heads, state_dim // 4), dtype=torch.float32)
47+ da_cs = torch.randn((batch, heads, seqlen), dtype=torch.float32) / 8
48+ da_cs_rev = torch.randn_like(da_cs) / 8
49+ dt = torch.rand_like(da_cs) / 10 + 0.05
50+ trap = torch.randn((batch, heads, seqlen), dtype=dtype)
51+ d = torch.randn((heads,), dtype=torch.float32) / 8 if has_d else None
52+ nchunks = (seqlen + chunk_size - 1) // chunk_size
53+ segsum = torch.randn((batch, heads, nchunks, chunk_size, chunk_size), dtype=torch.float32) / 8
54+ return (
55+ dout,
56+ q,
57+ k,
58+ v,
59+ q_bias,
60+ k_bias,
61+ mimo_v,
62+ mimo_o,
63+ z,
64+ mimo_z,
65+ angles,
66+ da_cs,
67+ da_cs_rev,
68+ dt,
69+ trap,
70+ d,
71+ segsum,
72+ chunk_size,
73+ 4,
74+ )
75+ 
76+ 
77+def _to_device(value, device):
78+ return value.to(device) if isinstance(value, torch.Tensor) else value
79+ 
80+ 
81+def _run_case(args):
82+ pytest.importorskip("triton")
83+ device = get_available_device()
84+ if device == "cpu":
85+ pytest.skip("mamba_mimo_bwd_fwd Triton requires an NPU")
86+ expected = mamba_mimo_bwd_fwd_golden(*args)
87+ actual = mamba_mimo_bwd_fwd(*tuple(_to_device(value, device) for value in args))
88+ for name, result, reference in zip(OUTPUT_NAMES, actual, expected):
89+ if reference is None:
90+ assert result is None
91+ continue
92+ assert result.dtype == reference.dtype
93+ error_ratio = get_err_ratio(reference.float(), result.detach().cpu().float())
94+ assert error_ratio < 0.001, f"{name} error_ratio={error_ratio}"
95+ 
96+ 
97+@pytest.mark.parametrize(
98+ ("reduce_o", "has_z", "has_d"),
99+ [
100+ (True, False, False),
101+ (True, True, True),
102+ (False, False, True),
103+ (False, True, False),
104+ ],
105+)
106+def test_mamba_mimo_bwd_fwd_branches(reduce_o, has_z, has_d):
107+ _run_case(_make_inputs(reduce_o=reduce_o, has_z=has_z, has_d=has_d))
108+ 
109+ 
110+def test_mamba_mimo_bwd_fwd_gqa():
111+ _run_case(_make_inputs(heads=2, groups=1, has_z=True, has_d=True))
112+ 
113+ 
114+@pytest.mark.parametrize("seed", [0, 1, 42])
115+def test_mamba_mimo_bwd_fwd_bf16_tail_seeds(seed):
116+ # These are the most sensitive configurations in the independent A40 check.
117+ _run_case(_make_inputs(seed=seed, seqlen=35, heads=2, has_z=True, has_d=True))
118+ 
119+ 
120+def test_mamba_mimo_bwd_fwd_batch_groups():
121+ _run_case(_make_inputs(batch=2, heads=4, groups=2, seqlen=20, has_z=True, has_d=True))
122+ 
123+ 
124+def test_mamba_mimo_bwd_fwd_non_power_of_two_fused_tail():
125+ _run_case(_make_inputs(seqlen=13, rank=3, chunk_size=6, has_z=True, has_d=True))
126+ 
127+ 
128+def test_mamba_mimo_bwd_fwd_tail_chunk():
129+ _run_case(_make_inputs(seqlen=20, has_z=True, has_d=True))
130+ 
131+ 
132+@pytest.mark.parametrize(
133+ ("state_dim", "value_dim", "rank", "chunk_size"),
134+ [
135+ (32, 64, 4, 16),
136+ (128, 32, 4, 16),
137+ (64, 128, 4, 16),
138+ (128, 64, 8, 8),
139+ (64, 64, 3, 6),
140+ (64, 64, 1, 64),
141+ ],
142+)
143+def test_mamba_mimo_bwd_fwd_shapes(state_dim, value_dim, rank, chunk_size):
144+ _run_case(
145+ _make_inputs(
146+ seqlen=chunk_size * 2,
147+ state_dim=state_dim,
148+ value_dim=value_dim,
149+ rank=rank,
150+ chunk_size=chunk_size,
151+ )
152+ )
153+ 
154+ 
155+def test_invalid_q_shape():
156+ args = list(_make_inputs(seqlen=16))
157+ args[1] = torch.randn((1, 16, 64), dtype=torch.bfloat16)
158+ with pytest.raises(ValueError, match="q must have shape"):
159+ mamba_mimo_bwd_fwd(*args)
160+ 
161+ 
162+def test_invalid_empty_batch():
163+ args = _make_inputs(batch=0, seqlen=16)
164+ with pytest.raises(ValueError, match="B, S, R and G must be positive"):
165+ mamba_mimo_bwd_fwd(*args)
166+ 
167+ 
168+def test_invalid_q_dtype():
169+ args = list(_make_inputs(seqlen=16))
170+ args[1] = args[1].to(torch.float16)
171+ with pytest.raises(TypeError, match="torch.bfloat16"):
172+ mamba_mimo_bwd_fwd(*args)