Star14
74
代码
代码
Issues30
Pull Requests33
Actions
项目讨论
Wiki
分析
Star14
74
已关闭
feat/fix/docs/style/refactor/adaptor/chore/test(backend): pr describe #98
SovLyn创建于 7月12日关闭于 7月24日
feat/fix/docs/style/refactor/adaptor/chore/test(backend): pr describe #98
已关闭
SovLyn创建于 7月12日关闭于 7月24日
从已删除 :master合入到Ascend/MindSpeed-Opsmaster
共 12 个文件变更+2324-0
Adocs/triton/mamba3_siso_bwd.md+215-0
@@ -0,0 +1,215 @@
1+# mamba3_siso_bwd Triton-Ascend 实践说明
2+ 
3+## 算子功能
4+ 
5+`mamba3_siso_bwd` 迁移 Mamba3 SISO backward 的三个核心算子：
6+ 
7+1. **`compute_dqktheta`** — 通过 rotary embeddings 和 biases 计算 dQ、dK、dAngles、dScale、dGamma、dQ_bias、dK_bias 梯度。
8+2. **`compute_ddt_dtrap_dinput_states`** — 从 dScale 和 dGamma 计算 dDT 和 dTrap 梯度，并可选地计算 input state 梯度。
9+ 
10+三个 Triton kernel 对应：
11+ 
12+| Kernel | 功能 | Grid |
13+| --- | --- | --- |
14+| `mamba3_siso_bwd_kernel_rotary_bias_angles` | d Rotary+Bias+Angles kernel | `(nchunks, batch)` |
15+| `mamba3_siso_bwd_kernel_dk_state_post` | dK_state post-process kernel (via atomic add) | `(nheads, batch)` |
16+| `mamba3_siso_bwd_kernel_ddt_dtrap_dinput_states` | dDT, dTrap, dInput States kernel | `(nheads, batch)` |
17+ 
18+输入布局（`compute_dqktheta`）：
19+ 
20+- `q/k`: `[B, S, Hq, K]`
21+- `scale/gamma/dqk`: `[B, H, S]`
22+- `q_bias/k_bias`: `[H, K]`
23+- `angles`: `[B, S, H, A]`
24+- `dq_in/dk_in`: `[B, S, H, K]`
25+ 
26+输入布局（`compute_ddt_dtrap_dinput_states`）：
27+ 
28+- `dscale/dgamma/dt/trap`: `[B, H, S]`
29+ 
30+约束：`H % Hq == 0`，`A <= K / 2`，`chunk_size` 为 64。
31+ 
32+## 交付文件
33+ 
34+| 文件 | 作用 |
35+| --- | --- |
36+| `mindspeed_ops/api/triton/mamba3_siso_bwd.py` | 公开 API，负责 arch32 guard、参数转发和返回结果。 |
37+| `mindspeed_ops/arch32/triton/mamba3/mamba3_siso_bwd_impl.py` | Triton-Ascend 生产实现，包含三个 kernel 和三个 host 函数。 |
38+| `mindspeed_ops/arch32/triton/mamba3/mamba3_utils.py` | Triton JIT helper，提供 rotary 小角度 sin/cos 多项式、sigmoid、silu。 |
39+| `tests/unit_tests/triton/test_mamba3_siso_bwd.py` | pytest 精度单测，覆盖 compute_dqktheta 和 compute_ddt_dtrap_dinput_states。 |
40+| `tests/atk_tests/triton/mamba3_siso_bwd/generate_mamba3_siso_bwd.py` | ATK case generator，从 PR 内参数空间生成 case。 |
41+| `tests/atk_tests/triton/mamba3_siso_bwd/mamba3_siso_bwd.yaml` | ATK 配置，声明输入、dtype、范围、baseline API 与 Triton API。 |
42+| `tests/atk_tests/triton/mamba3_siso_bwd/triton_mamba3_siso_bwd.py` | ATK wrapper，测试侧 torch reference 与 candidate API 注册。 |
43+| `tests/atk_tests/triton/mamba3_siso_bwd/reference_impl.py` | PR 内 torch reference，仅用于精度测试和自证脚本。 |
44+| `tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py` | 自包含自证报告脚本，包含文件范围、静态 no-fallback、live 精度。 |
45+| `docs/triton/mamba3_siso_bwd.md` | 本说明文档。 |
46+ 
47+## 迁移要点
48+ 
49+- 将 CUDA/PTX `cos.approx.f32`、`sin.approx.f32` 替换为 Triton JIT 内可编译的多项式近似（`cos_approx`/`sin_approx`）。
50+- 移除 `assert z.is_cuda` 等 CUDA 专属断言，适配 NPU 设备。
51+- 移除 `triton.Config` 中的 `maxnreg` 参数，因为 NPU Triton 不支持该参数。
52+- 将 3D tensor 操作（`tl.reshape`/`tl.split`/`tl.join`）替换为 even/odd 列索引方式，避免 NPU 编译错误（`cannot align 2 axis for memref.alloc`）。
53+- 在 rotary embedding 的 GQA reduction 中使用分开的 even/odd 累加器，避免 3D 内存对齐问题。
54+- `mamba3_siso_bwd_kernel_dk_state_post` 不使用 autotune（使用 `tl.atomic_add` 写入 dK 和 dAngles，autotune 会导致多次覆写）。
55+- arch35 当前 fail loudly；不会静默切到 reference、CPU、vendor whole-op 或其它 backend fallback。
56+ 
57+## 评审术语说明
58+ 
59+- `candidate`：本 PR 提交的待评审实现，也就是生产 API 最终调用的 Triton-Ascend kernel。
60+- `reference`：只用于测试对比的正确性基准。本 PR 内的 reference 是 `tests/atk_tests/triton/mamba3_siso_bwd/reference_impl.py`，生产 API 不引用它。
61+- `fallback`：候选算子失败或变慢时，生产路径偷偷改用 torch、torch_npu 高层等价算子、CPU、上游 reference、vendor whole-op 或其它 backend 的实现。本 PR 禁止生产路径 fallback。
62+- `whole-op`：一个库函数直接完成整个算子功能。生产路径可以使用 torch/Triton 做张量、device 和 launch 管理，但不能调用现成 whole-op 顶替本 Triton 实现。
63+- `arch32/arch35`：MindSpeed-Ops 内的架构目录。当前 PR 交付 arch32 路径；arch35 没有交付实现，所以会明确报不支持，而不是静默切换到其它实现。
64+- `ATK`：仓内已有的自动化测试配置体系。本 PR 的 ATK 文件用于生成和执行测试，不属于生产调用路径。
65+- `GQA`：Grouped Query Attention，即 `nheads` 是 `nheads_qk` 的整数倍，多个 head 共享同一组 Q/K。
66+ 
67+## 评审人自证命令
68+ 
69+在 MindSpeed-Ops 仓库根目录执行：
70+ 
71+```shell
72+python tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py \
73+ --output /tmp/mamba3_siso_bwd_self_check.md
74+```
75+ 
76+这个脚本只依赖 PR 内文件和当前 MindSpeed-Ops 运行环境，不读取外部评测目录。它会生成 Markdown 报告并返回退出码：通过返回 `0`，失败返回非零。
77+ 
78+快速检查可先跑两条代表性 case：
79+ 
80+```shell
81+python tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py \
82+ --quick \
83+ --output /tmp/mamba3_siso_bwd_quick_self_check.md
84+```
85+ 
86+单元测试命令：
87+ 
88+```shell
89+pytest -q tests/unit_tests/triton/test_mamba3_siso_bwd.py
90+```
91+ 
92+ATK 配置入口：
93+ 
94+```text
95+tests/atk_tests/triton/mamba3_siso_bwd/mamba3_siso_bwd.yaml
96+```
97+ 
98+## 已完成自证报告
99+ 
100+本节是当前 PR 内容在本机 Ascend NPU 上已经跑出的完整自证结果，不是运行说明的占位。
101+复核命令仍保留在上一节，供评审人需要时重新生成同格式报告。
102+ 
103+执行命令：
104+ 
105+```shell
106+python tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py \
107+ --output /tmp/mamba3_siso_bwd_self_check_full.md
108+```
109+ 
110+执行时间：2026-07-03 00:12:00 +08:00。执行结果：退出码 `0`。
111+ 
112+### 0. 运行与结论
113+ 
114+- overall: PASS
115+- mode: `full`
116+- static_no_fallback: PASS
117+- live_cases: PASS
118+- repo_root: `<MindSpeed-Ops repository root>`
119+- external_workspace_required: `no`
120+ 
121+### 1. 提交范围
122+ 
123+| 文件 | 类型 | 用途 | 审查项 | 存在 |
124+| --- | --- | --- | --- | --- |
125+| `mindspeed_ops/api/triton/mamba3_siso_bwd.py` | production api | 公开 Python API；做 arch32 guard、参数转发和结果返回。 | 不能实现 reference 计算，不能调用 torch/torch_npu/vendor whole-op fallback。 | yes |
126+| `mindspeed_ops/arch32/triton/mamba3/mamba3_siso_bwd_impl.py` | production kernel | Triton-Ascend kernel 和 host dispatch；包含 rotary_bias_angles、dk_state_post、ddt_dtrap_dinput_states 三个 kernel。 | 生产计算必须进入本文件内 Triton kernel；dispatch 只能按运行时 shape/dtype/metadata。 | yes |
127+| `mindspeed_ops/arch32/triton/mamba3/mamba3_utils.py` | production helper | Triton JIT helper：rotary 小角度 sin/cos 多项式、sigmoid、silu。 | 仅提供 JIT helper，不读取外部 reference、case、golden 或上游包。 | yes |
128+| `tests/unit_tests/triton/test_mamba3_siso_bwd.py` | unit test | pytest 精度单测，覆盖 compute_dqktheta 和 compute_ddt_dtrap_dinput_states 两条生产路径。 | 测试侧可调用 torch reference；生产入口不能反向依赖测试代码。 | yes |
129+| `tests/atk_tests/triton/mamba3_siso_bwd/generate_mamba3_siso_bwd.py` | ATK generator | ATK 用例生成器；从提交内参数空间生成 shape/dtype/attr。 | 不读取外部 task cases、answer key 或私有 workload 文件。 | yes |
130+| `tests/atk_tests/triton/mamba3_siso_bwd/mamba3_siso_bwd.yaml` | ATK config | ATK 算子配置，声明输入、dtype、取值范围、baseline API 与 Triton API。 | 配置中的有效范围要与本交付支持边界一致。 | yes |
131+| `tests/atk_tests/triton/mamba3_siso_bwd/triton_mamba3_siso_bwd.py` | ATK wrapper | ATK baseline/candidate wrapper；baseline 仅用于测试侧对比。 | candidate wrapper 必须调用 `mindspeed_ops.api.triton.mamba3_siso_bwd`。 | yes |
132+| `tests/atk_tests/triton/mamba3_siso_bwd/reference_impl.py` | test reference | 提交内 torch reference，用于 UT/ATK/自证脚本精度校验。 | 只能被测试和自证脚本引用，生产路径不能 import 它。 | yes |
133+| `tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py` | self-check script | 一键生成 Markdown/JSON 自证报告，包含文件范围、静态 no-fallback、live 精度。 | 只依赖 PR 内文件和当前 MindSpeed-Ops 运行环境；不读取外部评测目录。 | yes |
134+| `docs/triton/mamba3_siso_bwd.md` | documentation | 中文实践说明和审查说明，记录支持边界、验证命令和证据口径。 | 文档声明必须能由本脚本或仓内测试复核。 | yes |
135+ 
136+### 2. 无 fallback 静态检查
137+ 
138+允许的生产依赖边界：
139+ 
140+| import/text | 用途 |
141+| --- | --- |
142+| `torch` | tensor/runtime API for input metadata and device checks |
143+| `torch_npu` | runtime availability probe only; no torch_npu whole-op call is allowed |
144+| `triton` | Triton-Ascend JIT and launch API |
145+| `triton.language` | Triton language primitives |
146+ 
147+生产文件 import 扫描：
148+ 
149+- `mindspeed_ops/api/triton/mamba3_siso_bwd.py` imports: `__future__`, `typing`, `torch`, `mindspeed_ops.api.triton.utils`, `mindspeed_ops.utils`, `mindspeed_ops.arch32.triton.mamba3.mamba3_siso_bwd_impl`, `mindspeed_ops.arch32.triton.mamba3.mamba3_siso_bwd_impl`
150+- `mindspeed_ops/arch32/triton/mamba3/mamba3_siso_bwd_impl.py` imports: `typing`, `torch`, `torch_npu`, `triton`, `triton.language`, `mindspeed_ops.arch32.triton.mamba3.mamba3_utils`
151+- `mindspeed_ops/arch32/triton/mamba3/mamba3_utils.py` imports: `triton`, `triton.language`
152+ 
153+- PASS: production files do not import or reference torch_npu high-level ops, golden/reference code, upstream runtime packages, peer code, or other backend fallback.
154+ 
155+### 3. live 精度
156+ 
157+accuracy pass rule: `max_abs <= abs_tol OR relative_rmse <= relative_rmse_tol`, abs_tol=`0.001`, relative_rmse_tol=`0.001`
158+ 
159+| case | dtype | shape | tensor | max_abs | rel_rmse | MARE | MERE | RMSE | result |
160+| ---: | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | --- |
161+| 1 | float16 | `B1_S32_Hq1_H1_K64_CS64` | dq | 0.0000608861 | 0.0002071885 | 0.0004769560 | 0.0001783582 | 0.0000237140 | PASS |
162+| 1 | float16 | `B1_S32_Hq1_H1_K64_CS64` | dk | 0.0000610352 | 0.0002124990 | 0.0019322314 | 0.0001810303 | 0.0000192188 | PASS |
163+| 1 | float16 | `B1_S32_Hq1_H1_K64_CS64` | dangles | 0.0000001490 | 0.0000001062 | 0.0000044614 | 0.0000002889 | 0.0000000627 | PASS |
164+| 1 | float16 | `B1_S32_Hq1_H1_K64_CS64` | dscale | 0.0000001490 | 0.0000001185 | 0.0000053441 | 0.0000003990 | 0.0000000592 | PASS |
165+| 1 | float16 | `B1_S32_Hq1_H1_K64_CS64` | dgamma | 0.0000000075 | 0.0000000737 | 0.0000712978 | 0.0000001770 | 0.0000000018 | PASS |
166+| 2 | float16 | `B1_S64_Hq1_H2_K64_CS64` | dq | 0.0001219809 | 0.0002143254 | 0.0004760298 | 0.0001779745 | 0.0000349229 | PASS |
167+| 2 | float16 | `B1_S64_Hq1_H2_K64_CS64` | dk | 0.0001215041 | 0.0002045313 | 0.0023482803 | 0.0001753308 | 0.0000255590 | PASS |
168+| 2 | float16 | `B1_S64_Hq1_H2_K64_CS64` | dangles | 0.0000004768 | 0.0000001480 | 0.0000028696 | 0.0000002554 | 0.0000001300 | PASS |
169+| 2 | float16 | `B1_S64_Hq1_H2_K64_CS64` | dscale | 0.0000004768 | 0.0000001556 | 0.0000050795 | 0.0000002916 | 0.0000001024 | PASS |
170+| 2 | float16 | `B1_S64_Hq1_H2_K64_CS64` | dgamma | 0.0000000149 | 0.0000000756 | 0.0000484480 | 0.0000001264 | 0.0000000018 | PASS |
171+| 3 | float16 | `B2_S64_Hq2_H4_K64_CS32` | dq | 0.0001221001 | 0.0002118507 | 0.0047994906 | 0.0001779644 | 0.0000343607 | PASS |
172+| 3 | float16 | `B2_S64_Hq2_H4_K64_CS32` | dk | 0.0001219809 | 0.0002095905 | 0.0009686586 | 0.0001767072 | 0.0000262123 | PASS |
173+| 3 | float16 | `B2_S64_Hq2_H4_K64_CS32` | dangles | 0.0000007153 | 0.0000001452 | 0.0000493808 | 0.0000005875 | 0.0000001792 | PASS |
174+| 3 | float16 | `B2_S64_Hq2_H4_K64_CS32` | dscale | 0.0000004768 | 0.0000001503 | 0.0000156663 | 0.0000004670 | 0.0000001486 | PASS |
175+| 3 | float16 | `B2_S64_Hq2_H4_K64_CS32` | dgamma | 0.0000000149 | 0.0000000751 | 0.0004196171 | 0.0000001859 | 0.0000000019 | PASS |
176+| 5 | bfloat16 | `B1_S96_Hq1_H2_K64_CS64` | dq | 0.0009753108 | 0.0017151781 | 0.0038860423 | 0.0014231033 | 0.0002790484 | PASS |
177+| 5 | bfloat16 | `B1_S96_Hq1_H2_K64_CS64` | dk | 0.0009760261 | 0.0016716727 | 0.0038712594 | 0.0014198506 | 0.0002089125 | PASS |
178+| 5 | bfloat16 | `B1_S96_Hq1_H2_K64_CS64` | dangles | 0.0000004768 | 0.0000001638 | 0.0000111523 | 0.0000004928 | 0.0000001713 | PASS |
179+| 5 | bfloat16 | `B1_S96_Hq1_H2_K64_CS64` | dscale | 0.0000004172 | 0.0000001620 | 0.0000065263 | 0.0000003077 | 0.0000001403 | PASS |
180+| 5 | bfloat16 | `B1_S96_Hq1_H2_K64_CS64` | dgamma | 0.0000000149 | 0.0000000768 | 0.0000585953 | 0.0000001356 | 0.0000000019 | PASS |
181+| 6 | bfloat16 | `B2_S96_Hq1_H2_K64_CS64` | dq | 0.0009765029 | 0.0016663175 | 0.0038535503 | 0.0014037021 | 0.0002714253 | PASS |
182+| 6 | bfloat16 | `B2_S96_Hq1_H2_K64_CS64` | dk | 0.0009760857 | 0.0016691156 | 0.0038709389 | 0.0013973597 | 0.0002083580 | PASS |
183+| 6 | bfloat16 | `B2_S96_Hq1_H2_K64_CS64` | dangles | 0.0000007451 | 0.0000001430 | 0.0000181485 | 0.0000007031 | 0.0000002289 | PASS |
184+| 6 | bfloat16 | `B2_S96_Hq1_H2_K64_CS64` | dscale | 0.0000007153 | 0.0000001719 | 0.0000036239 | 0.0000004094 | 0.0000002161 | PASS |
185+| 6 | bfloat16 | `B2_S96_Hq1_H2_K64_CS64` | dgamma | 0.0000000149 | 0.0000000767 | 0.0009313169 | 0.0000002564 | 0.0000000019 | PASS |
186+| 7 | float16 | `B1_S128_Hq2_H4_K64_CS64` | dq | 0.0001220107 | 0.0002117362 | 0.0014773572 | 0.0001783168 | 0.0000342835 | PASS |
187+| 7 | float16 | `B1_S128_Hq2_H4_K64_CS64` | dk | 0.0001219511 | 0.0002058915 | 0.0038244731 | 0.0001760291 | 0.0000258792 | PASS |
188+| 7 | float16 | `B1_S128_Hq2_H4_K64_CS64` | dangles | 0.0000009537 | 0.0000001636 | 0.0004074033 | 0.0000020912 | 0.0000002190 | PASS |
189+| 7 | float16 | `B1_S128_Hq2_H4_K64_CS64` | dscale | 0.0000007153 | 0.0000001603 | 0.0000094319 | 0.0000003863 | 0.0000001688 | PASS |
190+| 7 | float16 | `B1_S128_Hq2_H4_K64_CS64` | dgamma | 0.0000000149 | 0.0000000776 | 0.0001984816 | 0.0000001512 | 0.0000000019 | PASS |
191+| 101 | float16 | `B1_H1_S32` | dDT | 0.0000000149 | 0.0000000297 | 0.0000001294 | 0.0000001294 | 0.0000000044 | PASS |
192+| 101 | float16 | `B1_H1_S32` | dTrap | 0.0000000002 | 0.0000000211 | 0.0000000990 | 0.0000000990 | 0.0000000001 | PASS |
193+| 102 | float16 | `B1_H2_S64` | dDT | 0.0000000298 | 0.0000000394 | 0.0000001195 | 0.0000001195 | 0.0000000056 | PASS |
194+| 102 | float16 | `B1_H2_S64` | dTrap | 0.0000000005 | 0.0000000545 | 0.0000001615 | 0.0000001615 | 0.0000000001 | PASS |
195+| 103 | float16 | `B2_H2_S64` | dDT | 0.0000000298 | 0.0000000368 | 0.0000001622 | 0.0000001622 | 0.0000000059 | PASS |
196+| 103 | float16 | `B2_H2_S64` | dTrap | 0.0000000009 | 0.0000000468 | 0.0000001557 | 0.0000001557 | 0.0000000001 | PASS |
197+| 104 | bfloat16 | `B1_H4_S96` | dDT | 0.0000000298 | 0.0000000433 | 0.0000001861 | 0.0000001861 | 0.0000000067 | PASS |
198+| 104 | bfloat16 | `B1_H4_S96` | dTrap | 0.0000000005 | 0.0000000391 | 0.0000002167 | 0.0000002167 | 0.0000000001 | PASS |
199+| 105 | float16 | `B2_H4_S128` | dDT | 0.0000000298 | 0.0000000390 | 0.0000001735 | 0.0000001735 | 0.0000000061 | PASS |
200+| 105 | float16 | `B2_H4_S128` | dTrap | 0.0000000009 | 0.0000000473 | 0.0000001871 | 0.0000001871 | 0.0000000001 | PASS |
201+ 
202+### 4. 已知边界
203+ 
204+- arch: arch32 production path; arch35 fails loudly
205+- operator_scope: dense non-varlen backward path (dqktheta, ddt_dtrap)
206+- unsupported_scope: varlen, state-return, full backward combined
207+- known_limitation: headdim_qk=128 triggers NPU UB overflow (loop unrolling doubles buffer usage); current PR covers headdim_qk=64 only
208+ 
209+## 支持边界
210+ 
211+- 支持 arch32 Triton-Ascend 路径。
212+- 支持 dense、non-varlen backward 路径（dqktheta、ddt_dtrap）。
213+- 当前 PR 不覆盖 varlen、state-return、full backward combined 路径。
214+- headdim_qk=128 在 NPU 上因 UB 溢出（loop unrolling 翻倍 buffer 占用）暂不支持，当前 PR 覆盖 headdim_qk=64。
215+- 测试侧 torch reference 只存在于 `tests/atk_tests/.../reference_impl.py`，生产 API 不引用测试代码。
Amindspeed_ops/api/triton/mamba3_siso_bwd.py+66-0
@@ -0,0 +1,66 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+# Copyright (c) 2025, Dao AI Lab, Goombalab
3+ 
4+from __future__ import annotations
5+ 
6+from typing import Optional, Tuple
7+ 
8+import torch
9+ 
10+from mindspeed_ops.api.triton.utils import input_guard
11+from mindspeed_ops.utils import is_arch35
12+ 
13+__all__ = ["compute_dqktheta", "compute_ddt_dtrap_dinput_states"]
14+ 
15+ 
16+@input_guard
17+def compute_dqktheta(
18+ q: torch.Tensor,
19+ k: torch.Tensor,
20+ scale: torch.Tensor,
21+ gamma: torch.Tensor,
22+ q_bias: torch.Tensor,
23+ k_bias: torch.Tensor,
24+ angles: torch.Tensor,
25+ dq_in: torch.Tensor,
26+ dk_in: torch.Tensor,
27+ dqk: torch.Tensor,
28+ d_ok_state: Optional[torch.Tensor] = None,
29+ chunk_size: int = 64,
30+ Cu_Seqlens: Optional[torch.Tensor] = None,
31+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
32+ """Compute gradients through rotary embeddings and biases for Mamba-3 backward pass."""
33+ if is_arch35():
34+ raise NotImplementedError("mamba3_siso_bwd is currently implemented for arch32 only")
35+ 
36+ from mindspeed_ops.arch32.triton.mamba3.mamba3_siso_bwd_impl import compute_dqktheta as _impl
37+ 
38+ return _impl(
39+ q=q, k=k, scale=scale, gamma=gamma, q_bias=q_bias, k_bias=k_bias,
40+ angles=angles, dq_in=dq_in, dk_in=dk_in, dqk=dqk,
41+ d_ok_state=d_ok_state, chunk_size=chunk_size, Cu_Seqlens=Cu_Seqlens,
42+ )
43+ 
44+ 
45+@input_guard
46+def compute_ddt_dtrap_dinput_states(
47+ dscale: torch.Tensor,
48+ dgamma: torch.Tensor,
49+ dt: torch.Tensor,
50+ trap: torch.Tensor,
51+ d_issm_state: Optional[torch.Tensor] = None,
52+ input_k_state: Optional[torch.Tensor] = None,
53+ input_v_state: Optional[torch.Tensor] = None,
54+ Cu_Seqlens: Optional[torch.Tensor] = None,
55+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
56+ """Compute dDT, dTrap from dScale/dGamma, and optionally input state gradients."""
57+ if is_arch35():
58+ raise NotImplementedError("mamba3_siso_bwd is currently implemented for arch32 only")
59+ 
60+ from mindspeed_ops.arch32.triton.mamba3.mamba3_siso_bwd_impl import compute_ddt_dtrap_dinput_states as _impl
61+ 
62+ return _impl(
63+ dscale=dscale, dgamma=dgamma, dt=dt, trap=trap,
64+ d_issm_state=d_issm_state, input_k_state=input_k_state,
65+ input_v_state=input_v_state, Cu_Seqlens=Cu_Seqlens,
66+ )
Amindspeed_ops/arch32/triton/mamba3/mamba3_siso_bwd_impl.py+789-0
@@ -0,0 +1,789 @@
1+"""
2+Mamba-3 Backward Pass Triton Kernels for Ascend NPU.
3+ 
4+Copyright (c) 2026, Dao AI Lab, Goombalab
5+Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
6+"""
7+ 
8+from typing import Optional, Tuple
9+ 
10+import torch
11+import torch_npu
12+ 
13+import triton
14+import triton.language as tl
15+from mindspeed_ops.arch32.triton.mamba3.mamba3_utils import cos_approx, sin_approx, sigmoid_approx
16+ 
17+ 
18+# =============================================================================
19+# d Rotary+Bias Kernel
20+# =============================================================================
21+ 
22+ 
23+@triton.autotune(
24+ configs=[
25+ triton.Config({"multibuffer": False}, num_stages=1, num_warps=2)
26+ ],
27+ key=["CHUNK_SIZE", "BLOCK_HEADDIM_QK", "HEADDIM_QK", "GQA_RATIO"]
28+)
29+@triton.jit
30+def mamba3_siso_bwd_kernel_rotary_bias_angles(
31+ # Input tensors
32+ Q, K, Scale, Gamma, Q_bias, K_bias, Angles, dQ_in, dK_in, dQK,
33+ # Output tensors
34+ dQ, dK, dAngles, dScale, dGamma, dQ_bias, dK_bias,
35+ # Strides for inputs -------------------------------------------------------
36+ # Q: (batch, seqlen, nheads_qk, headdim_qk)
37+ stride_q_batch, stride_q_seqlen, stride_q_head, stride_q_qkdim,
38+ # K: (batch, seqlen, nheads_qk, headdim_qk)
39+ stride_k_batch, stride_k_seqlen, stride_k_head, stride_k_qkdim,
40+ # Scale: (batch, nheads, seqlen)
41+ stride_scale_batch, stride_scale_head, stride_scale_seqlen,
42+ # Gamma: (batch, nheads, seqlen)
43+ stride_gamma_batch, stride_gamma_head, stride_gamma_seqlen,
44+ # Q_bias: (nheads, headdim_qk)
45+ stride_q_bias_head, stride_q_bias_qkdim,
46+ # K_bias: (nheads, headdim_qk)
47+ stride_k_bias_head, stride_k_bias_qkdim,
48+ # Angles: (batch, seqlen, nheads, headdim_qk/2)
49+ stride_angles_batch, stride_angles_seqlen, stride_angles_head, stride_angles_qkdim,
50+ # dQ_in: (batch, seqlen, nheads, headdim_qk)
51+ stride_dq_in_batch, stride_dq_in_seqlen, stride_dq_in_head, stride_dq_in_qkdim,
52+ # dK_in: (batch, seqlen, nheads, headdim_qk)
53+ stride_dk_in_batch, stride_dk_in_seqlen, stride_dk_in_head, stride_dk_in_qkdim,
54+ # dQK: (batch, nheads, seqlen)
55+ stride_dqk_batch, stride_dqk_head, stride_dqk_seqlen,
56+ # Strides for outputs ------------------------------------------------------
57+ # dQ: (batch, seqlen, nheads_qk, headdim_qk)
58+ stride_dq_batch, stride_dq_seqlen, stride_dq_head, stride_dq_qkdim,
59+ # dK: (batch, seqlen, nheads_qk, headdim_qk)
60+ stride_dk_batch, stride_dk_seqlen, stride_dk_head, stride_dk_qkdim,
61+ # dAngles: (batch, seqlen, nheads, headdim_qk/2)
62+ stride_dangles_batch, stride_dangles_seqlen, stride_dangles_head, stride_dangles_qkdim,
63+ # dScale: (batch, nheads, n_qk_blocks, seqlen)
64+ stride_dscale_batch, stride_dscale_head, stride_dscale_nqkchunks, stride_dscale_seqlen,
65+ # dGamma: (batch, nheads, n_qk_blocks, seqlen)
66+ stride_dgamma_batch, stride_dgamma_head, stride_dgamma_nqkchunks, stride_dgamma_seqlen,
67+ # dQ_bias: (batch, nchunks, nheads, headdim_qk)
68+ stride_dq_bias_batch, stride_dq_bias_nchunks, stride_dq_bias_head, stride_dq_bias_qkdim,
69+ # dK_bias: (batch, nchunks, nheads, headdim_qk)
70+ stride_dk_bias_batch, stride_dk_bias_nchunks, stride_dk_bias_head, stride_dk_bias_qkdim,
71+ # ---- sizes ----
72+ seqlen, nheads_qk, nheads, headdim_qk, headdim_angles,
73+ CHUNK_SIZE: tl.constexpr,
74+ HEADDIM_QK: tl.constexpr,
75+ BLOCK_HEADDIM_QK: tl.constexpr,
76+ GQA_RATIO: tl.constexpr,
77+):
78+ """
79+ Grid: (nchunks, batch)
80+ Each program processes one (batch, chunk) pair.
81+ """
82+ pid_nchunk = tl.program_id(0)
83+ pid_batch = tl.program_id(1)
84+ 
85+ # Base offsets for inputs
86+ q_offset_base = pid_batch * stride_q_batch
87+ k_offset_base = pid_batch * stride_k_batch
88+ scale_offset_base = pid_batch * stride_scale_batch
89+ gamma_offset_base = pid_batch * stride_gamma_batch
90+ angle_offset_base = pid_batch * stride_angles_batch
91+ dq_in_offset_base = pid_batch * stride_dq_in_batch
92+ dk_in_offset_base = pid_batch * stride_dk_in_batch
93+ dqk_offset_base = pid_batch * stride_dqk_batch
94+ 
95+ # Base offsets for outputs
96+ dq_offset_base = pid_batch * stride_dq_batch
97+ dk_offset_base = pid_batch * stride_dk_batch
98+ dangle_offset_base = pid_batch * stride_dangles_batch
99+ dscale_offset_base = pid_batch * stride_dscale_batch
100+ dgamma_offset_base = pid_batch * stride_dgamma_batch
101+ dq_bias_offset_base = pid_batch * stride_dq_bias_batch + pid_nchunk * stride_dq_bias_nchunks
102+ dk_bias_offset_base = pid_batch * stride_dk_bias_batch + pid_nchunk * stride_dk_bias_nchunks
103+ 
104+ num_nheads_qk = HEADDIM_QK // BLOCK_HEADDIM_QK
105+ for nhead_qk_id in range(num_nheads_qk):
106+ offs_s = tl.arange(0, CHUNK_SIZE) + pid_nchunk * CHUNK_SIZE
107+ offs_dr = tl.arange(0, BLOCK_HEADDIM_QK // 2) + nhead_qk_id * (BLOCK_HEADDIM_QK // 2)
108+ # Even/odd column offsets for rotary pairs (avoids tl.split/tl.join for NPU)
109+ offs_d_even = tl.arange(0, BLOCK_HEADDIM_QK // 2) * 2 + nhead_qk_id * BLOCK_HEADDIM_QK
110+ offs_d_odd = tl.arange(0, BLOCK_HEADDIM_QK // 2) * 2 + 1 + nhead_qk_id * BLOCK_HEADDIM_QK
111+ 
112+ for qk_head_idx in range(nheads_qk):
113+ q_offset = q_offset_base + qk_head_idx * stride_q_head
114+ k_offset = k_offset_base + qk_head_idx * stride_k_head
115+ # Even/odd pointers for Q, K (load directly to avoid slice indexing on NPU)
116+ q_ptrs_even = Q + q_offset + offs_s[:, None] * stride_q_seqlen + offs_d_even[None, :] * stride_q_qkdim
117+ q_ptrs_odd = Q + q_offset + offs_s[:, None] * stride_q_seqlen + offs_d_odd[None, :] * stride_q_qkdim
118+ k_ptrs_even = K + k_offset + offs_s[:, None] * stride_k_seqlen + offs_d_even[None, :] * stride_k_qkdim
119+ k_ptrs_odd = K + k_offset + offs_s[:, None] * stride_k_seqlen + offs_d_odd[None, :] * stride_k_qkdim
120+ 
121+ mask_2d_even = (offs_s[:, None] < seqlen) & (offs_d_even[None, :] < headdim_qk)
122+ mask_2d_odd = (offs_s[:, None] < seqlen) & (offs_d_odd[None, :] < headdim_qk)
123+ 
124+ # Separate even/odd accumulators for GQA reduction
125+ dq_acc_even = tl.zeros((CHUNK_SIZE, BLOCK_HEADDIM_QK // 2), dtype=tl.float32)
126+ dq_acc_odd = tl.zeros((CHUNK_SIZE, BLOCK_HEADDIM_QK // 2), dtype=tl.float32)
127+ dk_acc_even = tl.zeros((CHUNK_SIZE, BLOCK_HEADDIM_QK // 2), dtype=tl.float32)
128+ dk_acc_odd = tl.zeros((CHUNK_SIZE, BLOCK_HEADDIM_QK // 2), dtype=tl.float32)
129+ 
130+ for gqa_idx in range(GQA_RATIO):
131+ nhead_idx = qk_head_idx * GQA_RATIO + gqa_idx
132+ 
133+ # Bias for this head (even/odd columns loaded separately)
134+ q_bias_even = tl.load(
135+ Q_bias + nhead_idx * stride_q_bias_head + offs_d_even * stride_q_bias_qkdim,
136+ mask=offs_d_even < headdim_qk, other=0.0).to(tl.float32)
137+ q_bias_odd = tl.load(
138+ Q_bias + nhead_idx * stride_q_bias_head + offs_d_odd * stride_q_bias_qkdim,
139+ mask=offs_d_odd < headdim_qk, other=0.0).to(tl.float32)
140+ k_bias_even = tl.load(
141+ K_bias + nhead_idx * stride_k_bias_head + offs_d_even * stride_k_bias_qkdim,
142+ mask=offs_d_even < headdim_qk, other=0.0).to(tl.float32)
143+ k_bias_odd = tl.load(
144+ K_bias + nhead_idx * stride_k_bias_head + offs_d_odd * stride_k_bias_qkdim,
145+ mask=offs_d_odd < headdim_qk, other=0.0).to(tl.float32)
146+ 
147+ # Q + bias, K + bias (even/odd columns loaded directly)
148+ q0_even = tl.load(q_ptrs_even, mask=mask_2d_even, other=0.0)
149+ q0_odd = tl.load(q_ptrs_odd, mask=mask_2d_odd, other=0.0)
150+ k0_even = tl.load(k_ptrs_even, mask=mask_2d_even, other=0.0)
151+ k0_odd = tl.load(k_ptrs_odd, mask=mask_2d_odd, other=0.0)
152+ Q_wb_even = q0_even + q_bias_even[None, :]
153+ Q_wb_odd = q0_odd + q_bias_odd[None, :]
154+ K_wb_even = k0_even + k_bias_even[None, :]
155+ K_wb_odd = k0_odd + k_bias_odd[None, :]
156+ 
157+ # dQK
158+ dqk_offset = dqk_offset_base + nhead_idx * stride_dqk_head
159+ dqk = tl.load(dQK + dqk_offset + offs_s * stride_dqk_seqlen, mask=offs_s < seqlen, other=0.0)
160+ 
161+ # Scale, Gamma
162+ scale_offset = scale_offset_base + nhead_idx * stride_scale_head
163+ gamma_offset = gamma_offset_base + nhead_idx * stride_gamma_head
164+ scale = tl.load(Scale + scale_offset + offs_s * stride_scale_seqlen, mask=offs_s < seqlen, other=0.0).to(tl.float32)
165+ gamma = tl.load(Gamma + gamma_offset + offs_s * stride_gamma_seqlen, mask=offs_s < seqlen, other=0.0).to(tl.float32)
166+ 
167+ # Angles
168+ angle_offset = angle_offset_base + nhead_idx * stride_angles_head
169+ theta = tl.load(
170+ Angles + angle_offset + offs_s[:, None] * stride_angles_seqlen + offs_dr[None, :] * stride_angles_qkdim,
171+ mask=(offs_dr[None, :] < headdim_angles) & (offs_s[:, None] < seqlen),
172+ other=0.0).to(tl.float32)
173+ 
174+ # dQ_in, dK_in (even/odd columns loaded directly)
175+ dq_in_offset = dq_in_offset_base + nhead_idx * stride_dq_in_head
176+ dk_in_offset = dk_in_offset_base + nhead_idx * stride_dk_in_head
177+ dQ_in_even = tl.load(dQ_in + dq_in_offset + offs_s[:, None] * stride_dq_in_seqlen + offs_d_even[None, :] * stride_dq_in_qkdim,
178+ mask=mask_2d_even, other=0.0)
179+ dQ_in_odd = tl.load(dQ_in + dq_in_offset + offs_s[:, None] * stride_dq_in_seqlen + offs_d_odd[None, :] * stride_dq_in_qkdim,
180+ mask=mask_2d_odd, other=0.0)
181+ dK_in_even = tl.load(dK_in + dk_in_offset + offs_s[:, None] * stride_dk_in_seqlen + offs_d_even[None, :] * stride_dk_in_qkdim,
182+ mask=mask_2d_even, other=0.0)
183+ dK_in_odd = tl.load(dK_in + dk_in_offset + offs_s[:, None] * stride_dk_in_seqlen + offs_d_odd[None, :] * stride_dk_in_qkdim,
184+ mask=mask_2d_odd, other=0.0)
185+ 
186+ # dGamma = dQK * (Q_wbias · K_wbias), dot via even/odd sums
187+ QK_dot = tl.sum(Q_wb_even * K_wb_even, axis=1) + tl.sum(Q_wb_odd * K_wb_odd, axis=1)
188+ d_gamma = dqk * QK_dot
189+ dgamma_store_offset = dgamma_offset_base + nhead_idx * stride_dgamma_head
190+ tl.store(
191+ dGamma + dgamma_store_offset + offs_s * stride_dgamma_seqlen + nhead_qk_id * stride_dgamma_nqkchunks,
192+ d_gamma, mask=offs_s < seqlen)
193+ 
194+ # cos/sin
195+ cos_angle = cos_approx(theta.to(tl.float32))
196+ sin_angle = sin_approx(theta.to(tl.float32))
197+ 
198+ # dScale = sum(dK_in * K_rot) using even/odd
199+ K_rot_even = K_wb_even * cos_angle - K_wb_odd * sin_angle
200+ K_rot_odd = K_wb_even * sin_angle + K_wb_odd * cos_angle
201+ 
202+ dscale_val = tl.sum(dK_in_even * K_rot_even, axis=1) + tl.sum(dK_in_odd * K_rot_odd, axis=1)
203+ dscale_store_offset = dscale_offset_base + nhead_idx * stride_dscale_head
204+ tl.store(
205+ dScale + dscale_store_offset + offs_s * stride_dscale_seqlen + nhead_qk_id * stride_dscale_nqkchunks,
206+ dscale_val, mask=offs_s < seqlen)
207+ 
208+ # Inverse rotary on even/odd halves (scale applied directly)
209+ dK_in_scaled_even = dK_in_even * scale[:, None]
210+ dK_in_scaled_odd = dK_in_odd * scale[:, None]
211+ 
212+ dq0 = dQ_in_even * cos_angle + dQ_in_odd * sin_angle
213+ dq1 = -dQ_in_even * sin_angle + dQ_in_odd * cos_angle
214+ dk0 = dK_in_scaled_even * cos_angle + dK_in_scaled_odd * sin_angle
215+ dk1 = -dK_in_scaled_even * sin_angle + dK_in_scaled_odd * cos_angle
216+ 
217+ # Add dQK path
218+ dqk_scaled = (dqk * gamma)[:, None]
219+ dQ_pre_even = dq0 + dqk_scaled * K_wb_even
220+ dQ_pre_odd = dq1 + dqk_scaled * K_wb_odd
221+ dK_pre_even = dk0 + dqk_scaled * Q_wb_even
222+ dK_pre_odd = dk1 + dqk_scaled * Q_wb_odd
223+ 
224+ # Accumulate for GQA reduction
225+ dq_acc_even += dQ_pre_even
226+ dq_acc_odd += dQ_pre_odd
227+ dk_acc_even += dK_pre_even
228+ dk_acc_odd += dK_pre_odd
229+ 
230+ # Store bias gradients (even/odd separately)
231+ dq_bias_out_even = tl.sum(dQ_pre_even, axis=0)
232+ dq_bias_out_odd = tl.sum(dQ_pre_odd, axis=0)
233+ dk_bias_out_even = tl.sum(dK_pre_even, axis=0)
234+ dk_bias_out_odd = tl.sum(dK_pre_odd, axis=0)
235+ dq_bias_store_offset = dq_bias_offset_base + nhead_idx * stride_dq_bias_head
236+ dk_bias_store_offset = dk_bias_offset_base + nhead_idx * stride_dk_bias_head
237+ tl.store(dQ_bias + dq_bias_store_offset + offs_d_even * stride_dq_bias_qkdim, dq_bias_out_even, mask=offs_d_even < headdim_qk)
238+ tl.store(dQ_bias + dq_bias_store_offset + offs_d_odd * stride_dq_bias_qkdim, dq_bias_out_odd, mask=offs_d_odd < headdim_qk)
239+ tl.store(dK_bias + dk_bias_store_offset + offs_d_even * stride_dk_bias_qkdim, dk_bias_out_even, mask=offs_d_even < headdim_qk)
240+ tl.store(dK_bias + dk_bias_store_offset + offs_d_odd * stride_dk_bias_qkdim, dk_bias_out_odd, mask=offs_d_odd < headdim_qk)
241+ 
242+ # dAngles
243+ dtheta_q = dQ_in_even * (-Q_wb_even * sin_angle - Q_wb_odd * cos_angle) + dQ_in_odd * (Q_wb_even * cos_angle - Q_wb_odd * sin_angle)
244+ dtheta_k = dK_in_scaled_even * (-K_wb_even * sin_angle - K_wb_odd * cos_angle) + dK_in_scaled_odd * (K_wb_even * cos_angle - K_wb_odd * sin_angle)
245+ dtheta = dtheta_q + dtheta_k
246+ 
247+ dangle_store_offset = dangle_offset_base + nhead_idx * stride_dangles_head
248+ tl.store(
249+ dAngles + dangle_store_offset + offs_s[:, None] * stride_dangles_seqlen + offs_dr[None, :] * stride_dangles_qkdim,
250+ dtheta, mask=(offs_dr[None, :] < headdim_angles) & (offs_s[:, None] < seqlen))
251+ 
252+ # End of GQA group: store accumulated dQ, dK using even/odd offsets
253+ dq_off = dq_offset_base + qk_head_idx * stride_dq_head
254+ dk_off = dk_offset_base + qk_head_idx * stride_dk_head
255+ tl.store(dQ + dq_off + offs_s[:, None] * stride_dq_seqlen + offs_d_even[None, :] * stride_dq_qkdim,
256+ dq_acc_even, mask=(offs_s[:, None] < seqlen) & (offs_d_even[None, :] < headdim_qk))
257+ tl.store(dQ + dq_off + offs_s[:, None] * stride_dq_seqlen + offs_d_odd[None, :] * stride_dq_qkdim,
258+ dq_acc_odd, mask=(offs_s[:, None] < seqlen) & (offs_d_odd[None, :] < headdim_qk))
259+ tl.store(dK + dk_off + offs_s[:, None] * stride_dk_seqlen + offs_d_even[None, :] * stride_dk_qkdim,
260+ dk_acc_even, mask=(offs_s[:, None] < seqlen) & (offs_d_even[None, :] < headdim_qk))
261+ tl.store(dK + dk_off + offs_s[:, None] * stride_dk_seqlen + offs_d_odd[None, :] * stride_dk_qkdim,
262+ dk_acc_odd, mask=(offs_s[:, None] < seqlen) & (offs_d_odd[None, :] < headdim_qk))
263+ 
264+ 
265+# NOTE: Do not autotune this kernel. It overwrites dK, dK_bias, dAngles via atomic adds
266+# and autotuning will lead to multiple overwrites.
267+@triton.jit
268+def mamba3_siso_bwd_kernel_dk_state_post(
269+ # Inputs tensors
270+ dK_State, Angles, K, K_bias, Cu_Seqlens,
271+ # Outputs tensors
272+ dK, dK_bias, dAngles,
273+ # Strides for dK_State: (num_sequences, nheads, headdim_qk)
274+ stride_dk_state_batch, stride_dk_state_head, stride_dk_state_qkdim,
275+ # Strides for Angles: (batch, seqlen, nheads, headdim_angles)
276+ stride_angles_batch, stride_angles_seqlen, stride_angles_head, stride_angles_qkdim,
277+ # Strides for K: (batch, seqlen, nheads_qk, headdim_qk)
278+ stride_k_batch, stride_k_seqlen, stride_k_head, stride_k_qkdim,
279+ # Strides for K_bias: (nheads, headdim_qk)
280+ stride_k_bias_head, stride_k_bias_qkdim,
281+ # Strides for Cu_Seqlens: (num_sequences + 1,)
282+ stride_cu_seqlen,
283+ # Strides for dK: (batch, seqlen, nheads_qk, headdim_qk)
284+ stride_dk_batch, stride_dk_seqlen, stride_dk_head, stride_dk_qkdim,
285+ # Strides for dK_bias: (nheads, headdim_qk)
286+ stride_dk_bias_head, stride_dk_bias_qkdim,
287+ # Strides for dAngles: (batch, seqlen, nheads, headdim_angles)
288+ stride_dangles_batch, stride_dangles_seqlen, stride_dangles_head, stride_dangles_qkdim,
289+ # Dimensions
290+ seqlen, headdim_qk, headdim_angles,
291+ HEADDIM_QK: tl.constexpr,
292+ GQA_RATIO: tl.constexpr,
293+ IS_VARLEN: tl.constexpr,
294+):
295+ """
296+ Post-kernel for d_ok_state contributions.
297+ Grid: (nheads, batch)
298+ """
299+ pid_head = tl.program_id(0)
300+ pid_batch = tl.program_id(1)
301+ 
302+ if IS_VARLEN:
303+ pid_seq = pid_batch
304+ pid_batch = 0
305+ cu_seqlen = tl.load(Cu_Seqlens + (pid_seq + 1) * stride_cu_seqlen).to(tl.int32)
306+ last_pos = cu_seqlen - 1
307+ else:
308+ pid_seq = 0
309+ last_pos = seqlen - 1
310+ 
311+ qk_head_idx = pid_head // GQA_RATIO
312+ offs_dr = tl.arange(0, HEADDIM_QK // 2)
313+ # Even/odd offsets to split rotary pairs without tl.split/tl.join on NPU
314+ offs_d_even = offs_dr * 2
315+ offs_d_odd = offs_dr * 2 + 1
316+ 
317+ # Load dK_State even/odd columns
318+ dk_state_base = dK_State + (pid_batch + pid_seq) * stride_dk_state_batch + pid_head * stride_dk_state_head
319+ dk_state_r0 = tl.load(dk_state_base + offs_d_even * stride_dk_state_qkdim, mask=offs_d_even < headdim_qk, other=0.0).to(tl.float32)
320+ dk_state_r1 = tl.load(dk_state_base + offs_d_odd * stride_dk_state_qkdim, mask=offs_d_odd < headdim_qk, other=0.0).to(tl.float32)
321+ 
322+ # Load angles at last position
323+ angles_base = Angles + pid_batch * stride_angles_batch + last_pos * stride_angles_seqlen + pid_head * stride_angles_head
324+ angles_val = tl.load(angles_base + offs_dr * stride_angles_qkdim, mask=offs_dr < headdim_angles, other=0.0).to(tl.float32)
325+ 
326+ cos_ang = cos_approx(angles_val)
327+ sin_ang = sin_approx(angles_val)
328+ 
329+ # Inverse rotary: dk_rotated (even/odd halves)
330+ dk0 = dk_state_r0 * cos_ang + dk_state_r1 * sin_ang
331+ dk1 = -dk_state_r0 * sin_ang + dk_state_r1 * cos_ang
332+ 
333+ # 1. Accumulate to dK (GQA reduction via atomic) - even/odd columns
334+ dk_base = dK + pid_batch * stride_dk_batch + last_pos * stride_dk_seqlen + qk_head_idx * stride_dk_head
335+ tl.atomic_add(dk_base + offs_d_even * stride_dk_qkdim, dk0, mask=offs_d_even < headdim_qk)
336+ tl.atomic_add(dk_base + offs_d_odd * stride_dk_qkdim, dk1, mask=offs_d_odd < headdim_qk)
337+ 
338+ # 2. Accumulate to dK_bias (batch reduction via atomic) - even/odd columns
339+ dk_bias_base = dK_bias + pid_head * stride_dk_bias_head
340+ tl.atomic_add(dk_bias_base + offs_d_even * stride_dk_bias_qkdim, dk0, mask=offs_d_even < headdim_qk)
341+ tl.atomic_add(dk_bias_base + offs_d_odd * stride_dk_bias_qkdim, dk1, mask=offs_d_odd < headdim_qk)
342+ 
343+ # 3. Compute dAngles
344+ k_base = K + pid_batch * stride_k_batch + last_pos * stride_k_seqlen + qk_head_idx * stride_k_head
345+ k_r0 = tl.load(k_base + offs_d_even * stride_k_qkdim, mask=offs_d_even < headdim_qk, other=0.0).to(tl.float32)
346+ k_r1 = tl.load(k_base + offs_d_odd * stride_k_qkdim, mask=offs_d_odd < headdim_qk, other=0.0).to(tl.float32)
347+ 
348+ k_bias_base = K_bias + pid_head * stride_k_bias_head
349+ kb_r0 = tl.load(k_bias_base + offs_d_even * stride_k_bias_qkdim, mask=offs_d_even < headdim_qk, other=0.0).to(tl.float32)
350+ kb_r1 = tl.load(k_bias_base + offs_d_odd * stride_k_bias_qkdim, mask=offs_d_odd < headdim_qk, other=0.0).to(tl.float32)
351+ 
352+ K_wbias_r0 = k_r0 + kb_r0
353+ K_wbias_r1 = k_r1 + kb_r1
354+ 
355+ dtheta_k = (dk_state_r0 * (-K_wbias_r0 * sin_ang - K_wbias_r1 * cos_ang) +
356+ dk_state_r1 * (K_wbias_r0 * cos_ang - K_wbias_r1 * sin_ang))
357+ 
358+ da_base = dAngles + pid_batch * stride_dangles_batch + last_pos * stride_dangles_seqlen + pid_head * stride_dangles_head
359+ tl.atomic_add(da_base + offs_dr * stride_dangles_qkdim, dtheta_k, mask=offs_dr < headdim_angles)
360+ 
361+ 
362+# =============================================================================
363+# dDT, dTrap, and dInput States Kernel
364+# =============================================================================
365+@triton.autotune(
366+ configs=[
367+ triton.Config({"CHUNK_SIZE": 64, "multibuffer": False}, num_stages=1, num_warps=2)
368+ ],
369+ key=["HEADDIM_V", "HEADDIM_QK", "HAS_INPUT_STATE", "IS_VARLEN"]
370+)
371+@triton.jit
372+def mamba3_siso_bwd_kernel_ddt_dtrap_dinput_states(
373+ # Input tensors
374+ dScale, dGamma, DT, Trap,
375+ d_ISSM_State, Input_K_State, Input_V_State, Cu_Seqlens,
376+ # Output tensors
377+ dDT, dTrap,
378+ dInput_SSM_State, dInput_K_State, dInput_V_State,
379+ # Strides for dScale: (batch, nheads, seqlen)
380+ stride_dscale_batch, stride_dscale_head, stride_dscale_seqlen,
381+ # Strides for dGamma: (batch, nheads, seqlen)
382+ stride_dgamma_batch, stride_dgamma_head, stride_dgamma_seqlen,
383+ # Strides for DT: (batch, nheads, seqlen)
384+ stride_dt_batch, stride_dt_head, stride_dt_seqlen,
385+ # Strides for Trap: (batch, nheads, seqlen)
386+ stride_trap_batch, stride_trap_head, stride_trap_seqlen,
387+ # Strides for d_ISSM_State: (num_sequences, nheads, headdim_v, headdim_qk)
388+ stride_d_issm_state_batch, stride_d_issm_state_head, stride_d_issm_state_vdim, stride_d_issm_state_qkdim,
389+ # Strides for Input_K_State: (num_sequences, nheads, headdim_qk)
390+ stride_input_k_state_batch, stride_input_k_state_head, stride_input_k_state_qkdim,
391+ # Strides for Input_V_State: (num_sequences, nheads, headdim_v)
392+ stride_input_v_state_batch, stride_input_v_state_head, stride_input_v_state_vdim,
393+ # Stride for Cu_Seqlens
394+ stride_cu_seqlen,
395+ # Strides for dDT: (batch, nheads, seqlen)
396+ stride_ddt_batch, stride_ddt_head, stride_ddt_seqlen,
397+ # Strides for dTrap: (batch, nheads, seqlen)
398+ stride_dtrap_batch, stride_dtrap_head, stride_dtrap_seqlen,
399+ # Strides for dInput_SSM_State: (num_sequences, nheads, headdim_v, headdim_qk)
400+ stride_dinput_ssm_state_batch, stride_dinput_ssm_state_head, stride_dinput_ssm_state_vdim, stride_dinput_ssm_state_qkdim,
401+ # Strides for dInput_K_State: (num_sequences, nheads, headdim_qk)
402+ stride_dinput_k_state_batch, stride_dinput_k_state_head, stride_dinput_k_state_qkdim,
403+ # Strides for dInput_V_State: (num_sequences, nheads, headdim_v)
404+ stride_dinput_v_state_batch, stride_dinput_v_state_head, stride_dinput_v_state_vdim,
405+ # Dimensions
406+ seqlen, headdim_v, headdim_qk,
407+ # Compile-time constants
408+ CHUNK_SIZE: tl.constexpr,
409+ HEADDIM_V: tl.constexpr,
410+ HEADDIM_QK: tl.constexpr,
411+ HAS_INPUT_STATE: tl.constexpr,
412+ IS_VARLEN: tl.constexpr,
413+):
414+ """
415+ Backward kernel for computing dDT, dTrap, and input state gradients.
416+ """
417+ pid_head = tl.program_id(0)
418+ pid_batch = tl.program_id(1)
419+ 
420+ if IS_VARLEN:
421+ pid_seq = pid_batch
422+ pid_batch = 0
423+ cu_seqlen = tl.load(Cu_Seqlens + pid_seq * stride_cu_seqlen).to(tl.int32)
424+ cu_seqlen_next = tl.load(Cu_Seqlens + (pid_seq + 1) * stride_cu_seqlen).to(tl.int32)
425+ seqlen = cu_seqlen_next - cu_seqlen
426+ else:
427+ pid_seq = 0
428+ cu_seqlen = 0
429+ 
430+ # Pointer Offsets
431+ dscale_offset = pid_batch * stride_dscale_batch + pid_head * stride_dscale_head + IS_VARLEN * cu_seqlen * stride_dscale_seqlen
432+ dgamma_offset = pid_batch * stride_dgamma_batch + pid_head * stride_dgamma_head + IS_VARLEN * cu_seqlen * stride_dgamma_seqlen
433+ dt_offset = pid_batch * stride_dt_batch + pid_head * stride_dt_head + IS_VARLEN * cu_seqlen * stride_dt_seqlen
434+ trap_offset = pid_batch * stride_trap_batch + pid_head * stride_trap_head + IS_VARLEN * cu_seqlen * stride_trap_seqlen
435+ ddt_offset = pid_batch * stride_ddt_batch + pid_head * stride_ddt_head + IS_VARLEN * cu_seqlen * stride_ddt_seqlen
436+ dtrap_offset = pid_batch * stride_dtrap_batch + pid_head * stride_dtrap_head + IS_VARLEN * cu_seqlen * stride_dtrap_seqlen
437+ 
438+ # Part 1: dDT and dTrap
439+ num_chunks = tl.cdiv(seqlen, CHUNK_SIZE)
440+ 
441+ for chunk_idx in range(num_chunks):
442+ offs_s = chunk_idx * CHUNK_SIZE + tl.arange(0, CHUNK_SIZE)
443+ mask = offs_s < seqlen
444+ 
445+ dscale_t = tl.load(dScale + dscale_offset + offs_s * stride_dscale_seqlen, mask=mask, other=0.0)
446+ dgamma_t = tl.load(dGamma + dgamma_offset + offs_s * stride_dgamma_seqlen, mask=mask, other=0.0)
447+ trap_presig_t = tl.load(Trap + trap_offset + offs_s * stride_trap_seqlen, mask=mask, other=0.0).to(tl.float32)
448+ trap_t = sigmoid_approx(trap_presig_t)
449+ dt_t = tl.load(DT + dt_offset + offs_s * stride_dt_seqlen, mask=mask, other=0.0)
450+ 
451+ offs_s_prev = offs_s - 1
452+ mask_prev = (offs_s_prev >= 0) & (offs_s_prev < seqlen)
453+ dscale_prev = tl.load(
454+ dScale + dscale_offset + offs_s_prev * stride_dscale_seqlen,
455+ mask=mask_prev,
456+ other=0.0
457+ )
458+ 
459+ ddt_t = (dgamma_t + dscale_t) * trap_t + dscale_prev * (1.0 - trap_t)
460+ dtrap_t = (dgamma_t + dscale_t) * dt_t - dscale_prev * dt_t
461+ dtrap_presig_t = dtrap_t * trap_t * (1.0 - trap_t)
462+ 
463+ tl.store(dDT + ddt_offset + offs_s * stride_ddt_seqlen, ddt_t, mask=mask)
464+ tl.store(dTrap + dtrap_offset + offs_s * stride_dtrap_seqlen, dtrap_presig_t, mask=mask)
465+ 
466+ # Part 2: Input State Gradients
467+ if HAS_INPUT_STATE:
468+ d_issm_offset = (pid_batch + pid_seq) * stride_d_issm_state_batch + pid_head * stride_d_issm_state_head
469+ input_k_offset = (pid_batch + pid_seq) * stride_input_k_state_batch + pid_head * stride_input_k_state_head
470+ input_v_offset = (pid_batch + pid_seq) * stride_input_v_state_batch + pid_head * stride_input_v_state_head
471+ dinput_ssm_offset = (pid_batch + pid_seq) * stride_dinput_ssm_state_batch + pid_head * stride_dinput_ssm_state_head
472+ dinput_k_offset = (pid_batch + pid_seq) * stride_dinput_k_state_batch + pid_head * stride_dinput_k_state_head
473+ dinput_v_offset = (pid_batch + pid_seq) * stride_dinput_v_state_batch + pid_head * stride_dinput_v_state_head
474+ 
475+ dt_0 = tl.load(DT + dt_offset).to(tl.float32)
476+ trap_presig_0 = tl.load(Trap + trap_offset).to(tl.float32)
477+ trap_0 = sigmoid_approx(trap_presig_0)
478+ scalar = dt_0 * (1.0 - trap_0)
479+ 
480+ offs_v = tl.arange(0, HEADDIM_V)
481+ offs_qk = tl.arange(0, HEADDIM_QK)
482+ 
483+ input_k = tl.load(
484+ Input_K_State + input_k_offset + offs_qk * stride_input_k_state_qkdim,
485+ mask=offs_qk < headdim_qk, other=0.0).to(tl.float32)
486+ input_v = tl.load(
487+ Input_V_State + input_v_offset + offs_v * stride_input_v_state_vdim,
488+ mask=offs_v < headdim_v, other=0.0).to(tl.float32)
489+ 
490+ d_issm = tl.load(
491+ d_ISSM_State + d_issm_offset +
492+ offs_v[:, None] * stride_d_issm_state_vdim +
493+ offs_qk[None, :] * stride_d_issm_state_qkdim,
494+ mask=(offs_v[:, None] < headdim_v) & (offs_qk[None, :] < headdim_qk),
495+ other=0.0).to(tl.float32)
496+ 
497+ tl.store(
498+ dInput_SSM_State + dinput_ssm_offset +
499+ offs_v[:, None] * stride_dinput_ssm_state_vdim +
500+ offs_qk[None, :] * stride_dinput_ssm_state_qkdim,
501+ d_issm,
502+ mask=(offs_v[:, None] < headdim_v) & (offs_qk[None, :] < headdim_qk),
503+ )
504+ 
505+ outer_product = input_v[:, None] * input_k[None, :]
506+ d_scalar = tl.sum(d_issm * outer_product)
507+ 
508+ dinput_v = tl.sum(d_issm * input_k[None, :], axis=1) * scalar
509+ dinput_k = tl.sum(d_issm * input_v[:, None], axis=0) * scalar
510+ 
511+ tl.store(dInput_V_State + dinput_v_offset + offs_v * stride_dinput_v_state_vdim, dinput_v, mask=offs_v < headdim_v)
512+ tl.store(dInput_K_State + dinput_k_offset + offs_qk * stride_dinput_k_state_qkdim, dinput_k, mask=offs_qk < headdim_qk)
513+ 
514+ ddt_0_contrib = d_scalar * (1.0 - trap_0)
515+ dtrap_0_contrib = d_scalar * (-dt_0)
516+ dtrap_0_presig_contrib = dtrap_0_contrib * trap_0 * (1.0 - trap_0)
517+ 
518+ tl.atomic_add(dDT + ddt_offset, ddt_0_contrib)
519+ tl.atomic_add(dTrap + dtrap_offset, dtrap_0_presig_contrib)
520+ 
521+ 
522+# =============================================================================
523+# Host Functions
524+# =============================================================================
525+def compute_dqktheta(
526+ q: torch.Tensor,
527+ k: torch.Tensor,
528+ scale: torch.Tensor,
529+ gamma: torch.Tensor,
530+ q_bias: torch.Tensor,
531+ k_bias: torch.Tensor,
532+ angles: torch.Tensor,
533+ dq_in: torch.Tensor,
534+ dk_in: torch.Tensor,
535+ dqk: torch.Tensor,
536+ d_ok_state: Optional[torch.Tensor] = None,
537+ chunk_size: int = 64,
538+ Cu_Seqlens: Optional[torch.Tensor] = None,
539+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
540+ """Compute gradients through rotary embeddings and biases for Mamba-3 backward pass."""
541+ batch, seqlen, nheads_qk, headdim_qk = q.shape
542+ assert q.shape == k.shape
543+ 
544+ nheads = scale.shape[1]
545+ nchunks = triton.cdiv(seqlen, chunk_size)
546+ GQA_RATIO = nheads // nheads_qk
547+ 
548+ assert scale.shape == (batch, nheads, seqlen)
549+ assert gamma.shape == (batch, nheads, seqlen)
550+ assert q_bias.shape == (nheads, headdim_qk)
551+ assert k_bias.shape == (nheads, headdim_qk)
552+ headdim_angles = angles.shape[-1]
553+ assert angles.shape == (batch, seqlen, nheads, headdim_angles)
554+ assert dq_in.shape == (batch, seqlen, nheads, headdim_qk)
555+ assert dk_in.shape == (batch, seqlen, nheads, headdim_qk)
556+ assert dqk.shape == (batch, nheads, seqlen)
557+ if d_ok_state is not None:
558+ num_sequences = Cu_Seqlens.shape[0] - 1 if Cu_Seqlens is not None else batch
559+ assert d_ok_state.shape == (num_sequences, nheads, headdim_qk)
560+ assert nheads % nheads_qk == 0, "nheads must be multiple of nheads_qk for GQA support"
561+ 
562+ # Ensure contiguity
563+ q = q.contiguous() if not q.is_contiguous() else q
564+ k = k.contiguous() if not k.is_contiguous() else k
565+ scale = scale.contiguous() if not scale.is_contiguous() else scale
566+ gamma = gamma.contiguous() if not gamma.is_contiguous() else gamma
567+ dqk = dqk.contiguous() if not dqk.is_contiguous() else dqk
568+ angles = angles.contiguous() if not angles.is_contiguous() else angles
569+ dq_in = dq_in.contiguous() if not dq_in.is_contiguous() else dq_in
570+ dk_in = dk_in.contiguous() if not dk_in.is_contiguous() else dk_in
571+ q_bias = q_bias.contiguous() if q_bias.stride(-1) != 1 else q_bias
572+ k_bias = k_bias.contiguous() if k_bias.stride(-1) != 1 else k_bias
573+ if d_ok_state is not None and not d_ok_state.is_contiguous():
574+ d_ok_state = d_ok_state.contiguous()
575+ 
576+ HEADDIM_QK = triton.next_power_of_2(headdim_qk)
577+ BLOCK_HEADDIM_QK = min(HEADDIM_QK, 64)
578+ 
579+ dq = torch.empty((batch, seqlen, nheads_qk, headdim_qk), dtype=dq_in.dtype, device=q.device)
580+ dk = torch.empty((batch, seqlen, nheads_qk, headdim_qk), dtype=dk_in.dtype, device=k.device)
581+ dangles = torch.empty((batch, seqlen, nheads, headdim_angles), dtype=angles.dtype, device=angles.device)
582+ dscale = torch.empty((batch, nheads, HEADDIM_QK // BLOCK_HEADDIM_QK, seqlen), dtype=scale.dtype, device=scale.device)
583+ dgamma = torch.empty((batch, nheads, HEADDIM_QK // BLOCK_HEADDIM_QK, seqlen), dtype=gamma.dtype, device=gamma.device)
584+ dq_bias_partial = torch.empty((batch, nchunks, nheads, headdim_qk), dtype=torch.float32, device=q.device)
585+ dk_bias_partial = torch.empty((batch, nchunks, nheads, headdim_qk), dtype=torch.float32, device=k.device)
586+ 
587+ grid = (nchunks, batch)
588+ 
589+ mamba3_siso_bwd_kernel_rotary_bias_angles[grid](
590+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk,
591+ dq, dk, dangles, dscale, dgamma, dq_bias_partial, dk_bias_partial,
592+ q.stride(0), q.stride(1), q.stride(2), q.stride(3),
593+ k.stride(0), k.stride(1), k.stride(2), k.stride(3),
594+ scale.stride(0), scale.stride(1), scale.stride(2),
595+ gamma.stride(0), gamma.stride(1), gamma.stride(2),
596+ q_bias.stride(0), q_bias.stride(1),
597+ k_bias.stride(0), k_bias.stride(1),
598+ angles.stride(0), angles.stride(1), angles.stride(2), angles.stride(3),
599+ dq_in.stride(0), dq_in.stride(1), dq_in.stride(2), dq_in.stride(3),
600+ dk_in.stride(0), dk_in.stride(1), dk_in.stride(2), dk_in.stride(3),
601+ dqk.stride(0), dqk.stride(1), dqk.stride(2),
602+ dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3),
603+ dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
604+ dangles.stride(0), dangles.stride(1), dangles.stride(2), dangles.stride(3),
605+ dscale.stride(0), dscale.stride(1), dscale.stride(2), dscale.stride(3),
606+ dgamma.stride(0), dgamma.stride(1), dgamma.stride(2), dgamma.stride(3),
607+ dq_bias_partial.stride(0), dq_bias_partial.stride(1),
608+ dq_bias_partial.stride(2), dq_bias_partial.stride(3),
609+ dk_bias_partial.stride(0), dk_bias_partial.stride(1),
610+ dk_bias_partial.stride(2), dk_bias_partial.stride(3),
611+ seqlen, nheads_qk, nheads, headdim_qk, headdim_angles,
612+ CHUNK_SIZE=chunk_size,
613+ HEADDIM_QK=HEADDIM_QK,
614+ BLOCK_HEADDIM_QK=BLOCK_HEADDIM_QK,
615+ GQA_RATIO=GQA_RATIO,
616+ )
617+ 
618+ dscale = torch.sum(dscale, dim=2)
619+ dgamma = torch.sum(dgamma, dim=2)
620+ dq_bias = dq_bias_partial.sum(dim=(0, 1))
621+ dk_bias = dk_bias_partial.sum(dim=(0, 1))
622+ 
623+ if d_ok_state is not None:
624+ apply_dk_state_post(
625+ d_ok_state, angles, k, k_bias, dk, dk_bias, dangles, Cu_Seqlens
626+ )
627+ return dq, dk, dq_bias, dk_bias, dangles, dscale, dgamma
628+ 
629+ 
630+def apply_dk_state_post(
631+ d_ok_state: torch.Tensor,
632+ angles: torch.Tensor,
633+ k: torch.Tensor,
634+ k_bias: torch.Tensor,
635+ dk: torch.Tensor,
636+ dk_bias: torch.Tensor,
637+ dangles: torch.Tensor,
638+ Cu_Seqlens: Optional[torch.Tensor] = None,
639+):
640+ batch, seqlen, nheads, headdim_angles = angles.shape
641+ _, _, headdim_qk = d_ok_state.shape
642+ nheads_qk = k.shape[2]
643+ GQA_RATIO = nheads // nheads_qk
644+ 
645+ is_varlen = Cu_Seqlens is not None
646+ if is_varlen:
647+ num_sequences = Cu_Seqlens.shape[0] - 1
648+ assert batch == 1
649+ else:
650+ num_sequences = batch
651+ 
652+ d_ok_state = d_ok_state.contiguous() if not d_ok_state.is_contiguous() else d_ok_state
653+ angles = angles.contiguous() if not angles.is_contiguous() else angles
654+ k = k.contiguous() if not k.is_contiguous() else k
655+ k_bias = k_bias.contiguous() if not k_bias.is_contiguous() else k_bias
656+ 
657+ HEADDIM_QK = triton.next_power_of_2(headdim_qk)
658+ 
659+ grid = (nheads, num_sequences)
660+ 
661+ mamba3_siso_bwd_kernel_dk_state_post[grid](
662+ d_ok_state, angles, k, k_bias, Cu_Seqlens,
663+ dk, dk_bias, dangles,
664+ d_ok_state.stride(0), d_ok_state.stride(1), d_ok_state.stride(2),
665+ angles.stride(0), angles.stride(1), angles.stride(2), angles.stride(3),
666+ k.stride(0), k.stride(1), k.stride(2), k.stride(3),
667+ k_bias.stride(0), k_bias.stride(1),
668+ Cu_Seqlens.stride(0) if is_varlen else 0,
669+ dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
670+ dk_bias.stride(0), dk_bias.stride(1),
671+ dangles.stride(0), dangles.stride(1), dangles.stride(2), dangles.stride(3),
672+ seqlen, headdim_qk, headdim_angles,
673+ HEADDIM_QK=HEADDIM_QK,
674+ GQA_RATIO=GQA_RATIO,
675+ IS_VARLEN=is_varlen,
676+ num_warps=2,
677+ num_stages=3,
678+ )
679+ 
680+ 
681+def compute_ddt_dtrap_dinput_states(
682+ dscale: torch.Tensor,
683+ dgamma: torch.Tensor,
684+ dt: torch.Tensor,
685+ trap: torch.Tensor,
686+ d_issm_state: Optional[torch.Tensor] = None,
687+ input_k_state: Optional[torch.Tensor] = None,
688+ input_v_state: Optional[torch.Tensor] = None,
689+ Cu_Seqlens: Optional[torch.Tensor] = None,
690+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
691+ """Compute dDT, dTrap from dScale/dGamma, and optionally input state gradients."""
692+ batch, nheads, seqlen = dscale.shape
693+ has_input_state = d_issm_state is not None
694+ is_varlen = Cu_Seqlens is not None
695+ 
696+ if is_varlen:
697+ num_sequences = Cu_Seqlens.shape[0] - 1
698+ assert batch == 1
699+ else:
700+ num_sequences = batch
701+ 
702+ assert dgamma.shape == (batch, nheads, seqlen)
703+ assert dt.shape == (batch, nheads, seqlen)
704+ assert trap.shape == (batch, nheads, seqlen)
705+ 
706+ if has_input_state:
707+ assert input_k_state is not None and input_v_state is not None
708+ headdim_v, headdim_qk = d_issm_state.shape[2], d_issm_state.shape[3]
709+ assert d_issm_state.shape == (num_sequences, nheads, headdim_v, headdim_qk)
710+ assert input_k_state.shape == (num_sequences, nheads, headdim_qk)
711+ assert input_v_state.shape == (num_sequences, nheads, headdim_v)
712+ else:
713+ headdim_v, headdim_qk = 64, 128
714+ 
715+ dscale = dscale.contiguous() if not dscale.is_contiguous() else dscale
716+ dgamma = dgamma.contiguous() if not dgamma.is_contiguous() else dgamma
717+ dt = dt.contiguous() if not dt.is_contiguous() else dt
718+ trap = trap.contiguous() if not trap.is_contiguous() else trap
719+ 
720+ if has_input_state:
721+ d_issm_state = d_issm_state.contiguous() if not d_issm_state.is_contiguous() else d_issm_state
722+ input_k_state = input_k_state.contiguous() if not input_k_state.is_contiguous() else input_k_state
723+ input_v_state = input_v_state.contiguous() if not input_v_state.is_contiguous() else input_v_state
724+ 
725+ dDT = torch.empty_like(dt, dtype=torch.float32)
726+ dTrap = torch.empty_like(trap, dtype=torch.float32)
727+ 
728+ if has_input_state:
729+ d_Input_SSM_State = torch.empty_like(d_issm_state)
730+ d_Input_K_State = torch.empty((num_sequences, nheads, headdim_qk), dtype=torch.float32, device=dt.device)
731+ d_Input_V_State = torch.empty((num_sequences, nheads, headdim_v), dtype=torch.float32, device=dt.device)
732+ else:
733+ d_Input_SSM_State = None
734+ d_Input_K_State = None
735+ d_Input_V_State = None
736+ 
737+ HEADDIM_V = triton.next_power_of_2(headdim_v) if has_input_state else 64
738+ HEADDIM_QK = triton.next_power_of_2(headdim_qk) if has_input_state else 128
739+ 
740+ if is_varlen:
741+ grid = (nheads, num_sequences)
742+ else:
743+ grid = (nheads, batch)
744+ 
745+ mamba3_siso_bwd_kernel_ddt_dtrap_dinput_states[grid](
746+ dscale, dgamma, dt, trap,
747+ d_issm_state if has_input_state else dscale,
748+ input_k_state if has_input_state else dscale,
749+ input_v_state if has_input_state else dscale,
750+ Cu_Seqlens,
751+ dDT, dTrap,
752+ d_Input_SSM_State if has_input_state else dDT,
753+ d_Input_K_State if has_input_state else dDT,
754+ d_Input_V_State if has_input_state else dDT,
755+ dscale.stride(0), dscale.stride(1), dscale.stride(2),
756+ dgamma.stride(0), dgamma.stride(1), dgamma.stride(2),
757+ dt.stride(0), dt.stride(1), dt.stride(2),
758+ trap.stride(0), trap.stride(1), trap.stride(2),
759+ d_issm_state.stride(0) if has_input_state else 0,
760+ d_issm_state.stride(1) if has_input_state else 0,
761+ d_issm_state.stride(2) if has_input_state else 0,
762+ d_issm_state.stride(3) if has_input_state else 0,
763+ input_k_state.stride(0) if has_input_state else 0,
764+ input_k_state.stride(1) if has_input_state else 0,
765+ input_k_state.stride(2) if has_input_state else 0,
766+ input_v_state.stride(0) if has_input_state else 0,
767+ input_v_state.stride(1) if has_input_state else 0,
768+ input_v_state.stride(2) if has_input_state else 0,
769+ Cu_Seqlens.stride(0) if Cu_Seqlens is not None else 0,
770+ dDT.stride(0), dDT.stride(1), dDT.stride(2),
771+ dTrap.stride(0), dTrap.stride(1), dTrap.stride(2),
772+ d_Input_SSM_State.stride(0) if has_input_state else 0,
773+ d_Input_SSM_State.stride(1) if has_input_state else 0,
774+ d_Input_SSM_State.stride(2) if has_input_state else 0,
775+ d_Input_SSM_State.stride(3) if has_input_state else 0,
776+ d_Input_K_State.stride(0) if has_input_state else 0,
777+ d_Input_K_State.stride(1) if has_input_state else 0,
778+ d_Input_K_State.stride(2) if has_input_state else 0,
779+ d_Input_V_State.stride(0) if has_input_state else 0,
780+ d_Input_V_State.stride(1) if has_input_state else 0,
781+ d_Input_V_State.stride(2) if has_input_state else 0,
782+ seqlen, headdim_v, headdim_qk,
783+ HEADDIM_V=HEADDIM_V,
784+ HEADDIM_QK=HEADDIM_QK,
785+ HAS_INPUT_STATE=has_input_state,
786+ IS_VARLEN=is_varlen,
787+ )
788+ 
789+ return dDT, dTrap, d_Input_SSM_State, d_Input_K_State, d_Input_V_State
Amindspeed_ops/arch32/triton/mamba3/mamba3_utils.py+39-0
@@ -0,0 +1,39 @@
1+"""
2+Mamba-3 Util Functions.
3+ 
4+Copyright (c) 2025, Dao AI Lab, Goombalab
5+"""
6+ 
7+import triton
8+import triton.language as tl
9+ 
10+@triton.jit
11+def cos_approx(x):
12+ """
13+ Polynomial cosine approximation for the small rotary-angle range.
14+ """
15+ x2 = x * x
16+ return 1.0 - 0.5 * x2 + 0.041666666666666664 * x2 * x2 - 0.001388888888888889 * x2 * x2 * x2
17+ 
18+ 
19+@triton.jit
20+def sin_approx(x):
21+ """
22+ Polynomial sine approximation for the small rotary-angle range.
23+ """
24+ x2 = x * x
25+ return x * (1.0 - 0.16666666666666666 * x2 + 0.008333333333333333 * x2 * x2 - 0.0001984126984126984 * x2 * x2 * x2)
26+ 
27+@triton.jit
28+def sigmoid_approx(x):
29+ """
30+ Sigmoid helper kept as a JIT-callable wrapper for schedule readability.
31+ """
32+ return tl.sigmoid(x)
33+ 
34+@triton.jit
35+def silu(x):
36+ """
37+ SiLU (Swish) activation function: x * sigmoid(x).
38+ """
39+ return x*tl.sigmoid(x)
Amindspeed_ops/arch35/triton/mamba3/__init__.py+1-0
@@ -0,0 +1 @@
1+# arch35 mamba3 package
Amindspeed_ops/arch35/triton/mamba3/mamba3_siso_bwd_impl.py+9-0
@@ -0,0 +1,9 @@
1+"""
2+Mamba-3 Backward Pass Triton Kernels for Ascend NPU (arch35 stub).
3+ 
4+Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
5+Copyright (c) 2025, Dao AI Lab, Goombalab
6+"""
7+ 
8+# arch35 path is not yet implemented. The API layer will raise NotImplementedError
9+# when is_arch35() returns True, so this stub is a placeholder for future work.
Atests/atk_tests/triton/mamba3_siso_bwd/generate_mamba3_siso_bwd.py+117-0
@@ -0,0 +1,117 @@
1+import itertools
2+import random
3+ 
4+from atk.case_generator.generator.base_generator import CaseGenerator
5+from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
6+from atk.configs.case_config import CaseConfig
7+ 
8+ 
9+@GENERATOR_REGISTRY.register("generate_mamba3_siso_bwd_dqktheta")
10+class Mamba3SisoBwdDqkthetaGenerator(CaseGenerator):
11+ CASES = list(
12+ itertools.product(
13+ [1, 2], # batch
14+ [32, 64, 96, 128], # seqlen
15+ [1, 2], # q/k heads
16+ [1, 2, 4], # full heads
17+ [64, 128], # q/k dim
18+ [32, 64], # chunk size
19+ )
20+ )
21+ 
22+ _case_pool = None
23+ _case_index = 0
24+ 
25+ @classmethod
26+ def _build_pool(cls):
27+ return [case for case in cls.CASES if case[3] % case[2] == 0]
atomgit-bot
atomgit-botatomgit-bot7月12日

🟡 Medium Priority

changed line: generate_mamba3_siso_bwd.py 第 17 行 [64, 128] 以及第 27 行 _build_pool 仅过滤 nheads % nheads_qk == 0，未过滤 headdim_qk == 128。

affected behavior: Mamba3SisoBwdDqkthetaGenerator 的 CASES 和 _build_pool 会生成 headdim_qk=128 的用例。YAML 中 dtype_numbers=10 配合 2 种 dtype 共产生 20 个用例，按生成顺序前 20 个中有约一半是 headdim_qk=128。

failure mode: 文档（docs/triton/mamba3_siso_bwd.md 第 207 行及 self_check_mamba3_siso_bwd.py 第 306 行）明确声明 headdim_qk=128 在 NPU 上因 UB 溢出暂不支持，当前 PR 仅覆盖 headdim_qk=64。ATK 框架执行这些用例时，kernel 会因 NPU UB 溢出而崩溃或产生错误结果，导致 ATK 测试大量失败。

trigger condition: 当 ATK 框架按 mamba3_siso_bwd.yaml 配置运行 generate_mamba3_siso_bwd_dqktheta 生成器时必然触发。

suggested fix: 在 _build_pool 中增加 case[4] == 64 过滤条件，同步在 after_case_config 的 else 分支中将 random.choice([64, 128]) 改为 64。

建议：在 _build_pool 的列表推导中增加 case[4] == 64 过滤条件，并在 else 分支中将 headdim_qk 固定为 64。具体修改：第 27 行改为 return [case for case in cls.CASES if case[3] % case[2] == 0 and case[4] == 64]；第 42 行改为 headdim_qk = 64。

likedislike
不准确？
28+ 
29+ def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
30+ if Mamba3SisoBwdDqkthetaGenerator._case_pool is None:
31+ Mamba3SisoBwdDqkthetaGenerator._case_pool = self._build_pool()
32+ 
33+ if Mamba3SisoBwdDqkthetaGenerator._case_index < len(Mamba3SisoBwdDqkthetaGenerator._case_pool):
34+ batch, seqlen, nheads_qk, nheads, headdim_qk, chunk_size = (
35+ Mamba3SisoBwdDqkthetaGenerator._case_pool[Mamba3SisoBwdDqkthetaGenerator._case_index]
36+ )
37+ else:
38+ batch = random.choice([1, 2])
39+ seqlen = random.choice([32, 64, 96, 128, 160])
40+ nheads_qk = random.choice([1, 2, 4])
41+ nheads = random.choice([h for h in [1, 2, 4, 8] if h % nheads_qk == 0])
42+ headdim_qk = random.choice([64, 128])
atomgit-bot
atomgit-botatomgit-bot7月12日

🟡 Medium Priority

changed line: generate_mamba3_siso_bwd.py 第 42 行 headdim_qk = random.choice([64, 128])。

affected behavior: 当 _case_pool 耗尽后，after_case_config 的 else 分支通过随机采样生成用例，其中 headdim_qk 有 50% 概率选到 128。

failure mode: 与 _build_pool 中未过滤 headdim_qk=128 的问题相同——文档明确声明 headdim_qk=128 在 NPU 上因 UB 溢出暂不支持。随机生成的 headdim_qk=128 用例同样会导致 kernel 崩溃或错误结果。

suggested fix: 将第 42 行改为 headdim_qk = 64，与 _build_pool 的修复保持一致。

建议：将第 42 行改为 headdim_qk = 64，与文档中"当前 PR 覆盖 headdim_qk=64"的声明一致。

likedislike
不准确？
43+ chunk_size = random.choice([32, 64])
44+ 
45+ q_input = case_config.inputs[0]
46+ dtype = q_input.dtype
47+ headdim_angles = headdim_qk // 2
48+ 
49+ shapes = [
50+ [batch, seqlen, nheads_qk, headdim_qk],
51+ [batch, seqlen, nheads_qk, headdim_qk],
52+ [batch, nheads, seqlen],
53+ [batch, nheads, seqlen],
54+ [nheads, headdim_qk],
55+ [nheads, headdim_qk],
56+ [batch, seqlen, nheads, headdim_angles],
57+ [batch, seqlen, nheads, headdim_qk],
58+ [batch, seqlen, nheads, headdim_qk],
59+ [batch, nheads, seqlen],
60+ ]
61+ dtypes = [dtype, dtype, "fp32", "fp32", dtype, dtype, "fp32", dtype, dtype, "fp32"]
62+ 
63+ for input_cfg, shape, input_dtype in zip(case_config.inputs[:10], shapes, dtypes):
64+ if input_cfg.type == "tensor":
65+ input_cfg.shape = shape
66+ input_cfg.dtype = input_dtype
67+ 
68+ case_config.inputs[10].value = chunk_size
69+ Mamba3SisoBwdDqkthetaGenerator._case_index += 1
70+ return case_config
71+ 
72+ 
73+@GENERATOR_REGISTRY.register("generate_mamba3_siso_bwd_ddt")
74+class Mamba3SisoBwdDdtGenerator(CaseGenerator):
75+ CASES = list(
76+ itertools.product(
77+ [1, 2], # batch
78+ [1, 2, 4], # nheads
79+ [32, 64, 96, 128], # seqlen
80+ )
81+ )
82+ 
83+ _case_pool = None
84+ _case_index = 0
85+ 
86+ @classmethod
87+ def _build_pool(cls):
88+ return cls.CASES
89+ 
90+ def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
91+ if Mamba3SisoBwdDdtGenerator._case_pool is None:
92+ Mamba3SisoBwdDdtGenerator._case_pool = self._build_pool()
93+ 
94+ if Mamba3SisoBwdDdtGenerator._case_index < len(Mamba3SisoBwdDdtGenerator._case_pool):
95+ batch, nheads, seqlen = (
96+ Mamba3SisoBwdDdtGenerator._case_pool[Mamba3SisoBwdDdtGenerator._case_index]
97+ )
98+ else:
99+ batch = random.choice([1, 2])
100+ nheads = random.choice([1, 2, 4, 8])
101+ seqlen = random.choice([32, 64, 96, 128, 160])
102+ 
103+ shapes = [
104+ [batch, nheads, seqlen],
105+ [batch, nheads, seqlen],
106+ [batch, nheads, seqlen],
107+ [batch, nheads, seqlen],
108+ ]
109+ dtypes = ["fp32", "fp32", "fp32", "fp32"]
110+ 
111+ for input_cfg, shape, input_dtype in zip(case_config.inputs[:4], shapes, dtypes):
112+ if input_cfg.type == "tensor":
113+ input_cfg.shape = shape
114+ input_cfg.dtype = input_dtype
115+ 
116+ Mamba3SisoBwdDdtGenerator._case_index += 1
117+ return case_config
Atests/atk_tests/triton/mamba3_siso_bwd/mamba3_siso_bwd.yaml+163-0
@@ -0,0 +1,163 @@
1+api: pytorch
2+version: v1.0
3+name: torch_mamba3_siso_bwd_dqktheta.TorchMamba3SisoBwdDqkthetaApi
4+api_type: torch_mamba3_siso_bwd_dqktheta
5+triton_name: triton_mamba3_siso_bwd_dqktheta.TritonMamba3SisoBwdDqkthetaApi
6+triton_api_type: triton_mamba3_siso_bwd_dqktheta
7+dtype_numbers: 10
8+generate: generate_mamba3_siso_bwd_dqktheta
9+backward: false
10+standard:
11+ acc: single_bm
12+ perf: not_key
13+inputs:
14+ - name: q
15+ type: tensor
16+ required: true
17+ dtypes:
18+ values: [ fp16, bf16 ]
19+ ranges:
20+ valid:
21+ values: [ [-0.2, 0.2] ]
22+ invalid:
23+ values: [ [-0.2, 0.2] ]
24+ shapes:
25+ dim_numbers:
26+ values: [ 4 ]
27+ max_length: 4294967295
28+ - name: k
29+ type: tensor
30+ required: true
31+ dtypes:
32+ values: [ fp16, bf16 ]
33+ ranges:
34+ valid:
35+ values: [ [-0.2, 0.2] ]
36+ invalid:
37+ values: [ [-0.2, 0.2] ]
38+ shapes:
39+ dim_numbers:
40+ values: [ 4 ]
41+ max_length: 4294967295
42+ - name: scale
43+ type: tensor
44+ required: true
45+ dtypes:
46+ values: [ fp32 ]
47+ ranges:
48+ valid:
49+ values: [ [0.0, 1.0] ]
50+ invalid:
51+ values: [ [0.0, 1.0] ]
52+ shapes:
53+ dim_numbers:
54+ values: [ 3 ]
55+ max_length: 4294967295
56+ - name: gamma
57+ type: tensor
58+ required: true
59+ dtypes:
60+ values: [ fp32 ]
61+ ranges:
62+ valid:
63+ values: [ [0.0, 0.3] ]
64+ invalid:
65+ values: [ [0.0, 0.3] ]
66+ shapes:
67+ dim_numbers:
68+ values: [ 3 ]
69+ max_length: 4294967295
70+ - name: q_bias
71+ type: tensor
72+ required: true
73+ dtypes:
74+ values: [ fp16, bf16 ]
75+ ranges:
76+ valid:
77+ values: [ [-0.05, 0.05] ]
78+ invalid:
79+ values: [ [-0.05, 0.05] ]
80+ shapes:
81+ dim_numbers:
82+ values: [ 2 ]
83+ max_length: 4294967295
84+ - name: k_bias
85+ type: tensor
86+ required: true
87+ dtypes:
88+ values: [ fp16, bf16 ]
89+ ranges:
90+ valid:
91+ values: [ [-0.05, 0.05] ]
92+ invalid:
93+ values: [ [-0.05, 0.05] ]
94+ shapes:
95+ dim_numbers:
96+ values: [ 2 ]
97+ max_length: 4294967295
98+ - name: angles
99+ type: tensor
100+ required: true
101+ dtypes:
102+ values: [ fp32 ]
103+ ranges:
104+ valid:
105+ values: [ [0.0, 0.25] ]
106+ invalid:
107+ values: [ [0.0, 0.25] ]
108+ shapes:
109+ dim_numbers:
110+ values: [ 4 ]
111+ max_length: 4294967295
112+ - name: dq_in
113+ type: tensor
114+ required: true
115+ dtypes:
116+ values: [ fp16, bf16 ]
117+ ranges:
118+ valid:
119+ values: [ [-0.2, 0.2] ]
120+ invalid:
121+ values: [ [-0.2, 0.2] ]
122+ shapes:
123+ dim_numbers:
124+ values: [ 4 ]
125+ max_length: 4294967295
126+ - name: dk_in
127+ type: tensor
128+ required: true
129+ dtypes:
130+ values: [ fp16, bf16 ]
131+ ranges:
132+ valid:
133+ values: [ [-0.2, 0.2] ]
134+ invalid:
135+ values: [ [-0.2, 0.2] ]
136+ shapes:
137+ dim_numbers:
138+ values: [ 4 ]
139+ max_length: 4294967295
140+ - name: dqk
141+ type: tensor
142+ required: true
143+ dtypes:
144+ values: [ fp32 ]
145+ ranges:
146+ valid:
147+ values: [ [0.0, 0.2] ]
148+ invalid:
149+ values: [ [0.0, 0.2] ]
150+ shapes:
151+ dim_numbers:
152+ values: [ 3 ]
153+ max_length: 4294967295
154+ - name: chunk_size
155+ type: attr
156+ required: false
157+ dtypes:
158+ values: [ int ]
159+ ranges:
160+ valid:
161+ values: [ 32, 64 ]
162+ invalid:
163+ values: [ 32, 64 ]
atomgit-bot
atomgit-botatomgit-bot7月12日

🟡 Medium Priority

变更内容：mamba3_siso_bwd.yaml 仅配置了 torch_mamba3_siso_bwd_dqktheta / triton_mamba3_siso_bwd_dqktheta 测试用例（通过 name、triton_name、generate 字段），但 triton_mamba3_siso_bwd.py 中还注册了 torch_mamba3_siso_bwd_ddt 和 triton_mamba3_siso_bwd_ddt 两个测试类，且 generate_mamba3_siso_bwd.py 中已实现对应的 generate_mamba3_siso_bwd_ddt 生成器。缺少对应的 YAML 配置意味着 ATK 框架无法通过 YAML 驱动的方式运行 ddt 测试，该路径的 ATK 测试覆盖为空。虽然 ddt 路径在单元测试和 self_check 脚本中有覆盖，但 ATK 框架的 YAML 驱动测试是项目的标准测试流程。

建议：为 torch_mamba3_siso_bwd_ddt / triton_mamba3_siso_bwd_ddt 添加对应的 YAML 配置文件（可新建 mamba3_siso_bwd_ddt.yaml 或在现有 YAML 中增加 ddt 测试段），引用 generate_mamba3_siso_bwd_ddt 生成器，并定义 dscale、dgamma、dt、trap 四个输入张量的规格。

likedislike
不准确？
Atests/atk_tests/triton/mamba3_siso_bwd/reference_impl.py+187-0
@@ -0,0 +1,187 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+# Copyright (c) 2025, Dao AI Lab, Goombalab
3+ 
4+import torch
5+ 
6+ 
7+def _apply_rotary(x: torch.Tensor, angles: torch.Tensor) -> torch.Tensor:
8+ """Apply rotary embedding to x using angles."""
9+ out = x.to(torch.float32).clone()
10+ rotary_dim = min(angles.shape[-1], out.shape[-1] // 2)
11+ if rotary_dim <= 0:
12+ return out
13+ cos = torch.cos(angles[..., :rotary_dim].to(torch.float32))
14+ sin = torch.sin(angles[..., :rotary_dim].to(torch.float32))
15+ even = out[..., 0:2 * rotary_dim:2].clone()
16+ odd = out[..., 1:2 * rotary_dim:2].clone()
17+ out[..., 0:2 * rotary_dim:2] = even * cos - odd * sin
18+ out[..., 1:2 * rotary_dim:2] = even * sin + odd * cos
19+ return out
20+ 
21+ 
22+def _inv_apply_rotary(x: torch.Tensor, angles: torch.Tensor) -> torch.Tensor:
23+ """Apply inverse rotary embedding to x using angles."""
24+ out = x.to(torch.float32).clone()
25+ rotary_dim = min(angles.shape[-1], out.shape[-1] // 2)
26+ if rotary_dim <= 0:
27+ return out
28+ cos = torch.cos(angles[..., :rotary_dim].to(torch.float32))
29+ sin = torch.sin(angles[..., :rotary_dim].to(torch.float32))
30+ even = out[..., 0:2 * rotary_dim:2].clone()
31+ odd = out[..., 1:2 * rotary_dim:2].clone()
32+ out[..., 0:2 * rotary_dim:2] = even * cos + odd * sin
33+ out[..., 1:2 * rotary_dim:2] = -even * sin + odd * cos
34+ return out
35+ 
36+ 
37+def ref_compute_dqktheta(
38+ q: torch.Tensor,
39+ k: torch.Tensor,
40+ scale: torch.Tensor,
41+ gamma: torch.Tensor,
42+ q_bias: torch.Tensor,
43+ k_bias: torch.Tensor,
44+ angles: torch.Tensor,
45+ dq_in: torch.Tensor,
46+ dk_in: torch.Tensor,
47+ dqk: torch.Tensor,
48+ chunk_size: int = 64,
49+) -> tuple:
50+ """Reference implementation of compute_dqktheta."""
51+ batch, seqlen, nheads_qk, headdim_qk = q.shape
52+ nheads = scale.shape[1]
53+ GQA_RATIO = nheads // nheads_qk
54+ work_dtype = torch.float32
55+ 
56+ dq = torch.zeros((batch, seqlen, nheads_qk, headdim_qk), device=q.device, dtype=work_dtype)
57+ dk = torch.zeros((batch, seqlen, nheads_qk, headdim_qk), device=k.device, dtype=work_dtype)
58+ dangles = torch.zeros((batch, seqlen, nheads, angles.shape[-1]), device=angles.device, dtype=work_dtype)
59+ dscale = torch.zeros((batch, nheads, seqlen), device=scale.device, dtype=work_dtype)
60+ dgamma = torch.zeros((batch, nheads, seqlen), device=gamma.device, dtype=work_dtype)
61+ dq_bias = torch.zeros((nheads, headdim_qk), device=q.device, dtype=work_dtype)
62+ dk_bias = torch.zeros((nheads, headdim_qk), device=k.device, dtype=work_dtype)
63+ 
64+ for b_idx in range(batch):
65+ for qk_head_idx in range(nheads_qk):
66+ for gqa_idx in range(GQA_RATIO):
67+ nhead_idx = qk_head_idx * GQA_RATIO + gqa_idx
68+ 
69+ q_pre = q[b_idx, :, qk_head_idx, :].to(work_dtype) + q_bias[nhead_idx].to(work_dtype)
70+ k_pre = k[b_idx, :, qk_head_idx, :].to(work_dtype) + k_bias[nhead_idx].to(work_dtype)
71+ dq_in_h = dq_in[b_idx, :, nhead_idx, :].to(work_dtype)
72+ dk_in_h = dk_in[b_idx, :, nhead_idx, :].to(work_dtype)
73+ dqk_h = dqk[b_idx, nhead_idx, :].to(work_dtype)
74+ scale_h = scale[b_idx, nhead_idx, :].to(work_dtype)
75+ gamma_h = gamma[b_idx, nhead_idx, :].to(work_dtype)
76+ angles_h = angles[b_idx, :, nhead_idx, :].to(work_dtype)
77+ 
78+ # dGamma = dQK * (Q_wbias . K_wbias)
79+ qk_dot = (q_pre * k_pre).sum(dim=-1)
80+ dgamma[b_idx, nhead_idx, :] = dqk_h * qk_dot
81+ 
82+ # cos/sin
83+ cos_ang = torch.cos(angles_h)
84+ sin_ang = torch.sin(angles_h)
85+ 
86+ # dScale = sum(dK_in * K_rot)
87+ k_rot = _apply_rotary(k_pre, angles_h)
88+ dscale[b_idx, nhead_idx, :] = (dk_in_h * k_rot).sum(dim=-1)
89+ 
90+ # Inverse rotary
91+ dk_in_scaled = dk_in_h * scale_h.unsqueeze(-1)
92+ dq_unrot = _inv_apply_rotary(dq_in_h, angles_h)
93+ dk_unrot = _inv_apply_rotary(dk_in_scaled, angles_h)
94+ 
95+ # Add dQK path
96+ dqk_scaled = (dqk_h * gamma_h).unsqueeze(-1)
97+ dq_pre = dq_unrot + dqk_scaled * k_pre
98+ dk_pre = dk_unrot + dqk_scaled * q_pre
99+ 
100+ # Accumulate GQA
101+ dq[b_idx, :, qk_head_idx, :] += dq_pre
102+ dk[b_idx, :, qk_head_idx, :] += dk_pre
103+ dq_bias[nhead_idx, :] += dq_pre.sum(dim=0)
104+ dk_bias[nhead_idx, :] += dk_pre.sum(dim=0)
105+ 
106+ # dAngles
107+ rotary_dim = min(angles_h.shape[-1], headdim_qk // 2)
108+ if rotary_dim > 0:
109+ even_q = q_pre[:, 0:2 * rotary_dim:2]
110+ odd_q = q_pre[:, 1:2 * rotary_dim:2]
111+ even_k = k_pre[:, 0:2 * rotary_dim:2]
112+ odd_k = k_pre[:, 1:2 * rotary_dim:2]
113+ even_dq = dq_in_h[:, 0:2 * rotary_dim:2]
114+ odd_dq = dq_in_h[:, 1:2 * rotary_dim:2]
115+ even_dk = dk_in_scaled[:, 0:2 * rotary_dim:2]
116+ odd_dk = dk_in_scaled[:, 1:2 * rotary_dim:2]
117+ 
118+ dtheta_q = even_dq * (-even_q * sin_ang - odd_q * cos_ang) + \
119+ odd_dq * (even_q * cos_ang - odd_q * sin_ang)
120+ dtheta_k = even_dk * (-even_k * sin_ang - odd_k * cos_ang) + \
121+ odd_dk * (even_k * cos_ang - odd_k * sin_ang)
122+ dangles[b_idx, :, nhead_idx, :rotary_dim] = dtheta_q + dtheta_k
123+ 
124+ return dq, dk, dq_bias, dk_bias, dangles, dscale, dgamma
125+ 
126+ 
127+def ref_compute_ddt_dtrap_dinput_states(
128+ dscale: torch.Tensor,
129+ dgamma: torch.Tensor,
130+ dt: torch.Tensor,
131+ trap: torch.Tensor,
132+) -> tuple:
133+ """Reference implementation of compute_ddt_dtrap_dinput_states."""
134+ batch, nheads, seqlen = dscale.shape
135+ work_dtype = torch.float32
136+ 
137+ dDT = torch.zeros_like(dt, dtype=work_dtype)
138+ dTrap = torch.zeros_like(trap, dtype=work_dtype)
139+ 
140+ for b_idx in range(batch):
141+ for h_idx in range(nheads):
142+ dscale_h = dscale[b_idx, h_idx].to(work_dtype)
143+ dgamma_h = dgamma[b_idx, h_idx].to(work_dtype)
144+ dt_h = dt[b_idx, h_idx].to(work_dtype)
145+ trap_h = torch.sigmoid(trap[b_idx, h_idx].to(work_dtype))
146+ 
147+ for t in range(seqlen):
148+ dscale_prev = dscale_h[t - 1] if t > 0 else torch.tensor(0.0, device=dt.device, dtype=work_dtype)
149+ ddt_t = (dgamma_h[t] + dscale_h[t]) * trap_h[t] + dscale_prev * (1.0 - trap_h[t])
150+ dtrap_t = (dgamma_h[t] + dscale_h[t]) * dt_h[t] - dscale_prev * dt_h[t]
151+ dtrap_presig_t = dtrap_t * trap_h[t] * (1.0 - trap_h[t])
152+ 
153+ dDT[b_idx, h_idx, t] = ddt_t
154+ dTrap[b_idx, h_idx, t] = dtrap_presig_t
155+ 
156+ return dDT, dTrap
157+ 
158+ 
159+def make_inputs_dqktheta(batch, seqlen, nheads_qk, nheads, headdim_qk, dtype, device):
160+ """Create test inputs for compute_dqktheta."""
161+ torch.manual_seed(2026 + batch + seqlen + nheads + headdim_qk)
162+ headdim_angles = headdim_qk // 2
163+ 
164+ q = (torch.rand(batch, seqlen, nheads_qk, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
165+ k = (torch.rand(batch, seqlen, nheads_qk, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
166+ scale = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.5 + 0.5
167+ gamma = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.3
168+ q_bias = (torch.rand(nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.1
169+ k_bias = (torch.rand(nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.1
170+ angles = torch.rand(batch, seqlen, nheads, headdim_angles, device=device, dtype=torch.float32) * 0.25
171+ dq_in = (torch.rand(batch, seqlen, nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
172+ dk_in = (torch.rand(batch, seqlen, nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
173+ dqk = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.2
174+ 
175+ return q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk
176+ 
177+ 
178+def make_inputs_ddt(batch, nheads, seqlen, dtype, device):
179+ """Create test inputs for compute_ddt_dtrap_dinput_states."""
180+ torch.manual_seed(2027 + batch + seqlen + nheads)
181+ 
182+ dscale = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.2
183+ dgamma = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.2
184+ dt = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.1
185+ trap = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 2.0 - 1.0
186+ 
187+ return dscale, dgamma, dt, trap
Atests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py+401-0
@@ -0,0 +1,401 @@
1+#!/usr/bin/env python3
2+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
3+ 
4+from __future__ import annotations
5+ 
6+import argparse
7+import ast
8+import json
9+import math
10+import os
11+import statistics
12+from pathlib import Path
13+from typing import Any
14+ 
15+import torch
16+ 
17+from mindspeed_ops.api.triton.mamba3_siso_bwd import compute_dqktheta, compute_ddt_dtrap_dinput_states
18+from tests.atk_tests.triton.mamba3_siso_bwd.reference_impl import (
19+ make_inputs_dqktheta,
20+ make_inputs_ddt,
21+ ref_compute_dqktheta,
22+ ref_compute_ddt_dtrap_dinput_states,
23+)
24+from tests.utils import get_mare, get_mere, get_rmse
25+ 
26+ 
27+REPO_ROOT = Path(__file__).resolve().parents[4]
28+ 
29+FILE_INVENTORY = [
30+ {
31+ "path": "mindspeed_ops/api/triton/mamba3_siso_bwd.py",
32+ "kind": "production api",
33+ "purpose": "公开 Python API；做 arch32 guard、参数转发和结果返回。",
34+ "review_check": "不能实现 reference 计算，不能调用 torch/torch_npu/vendor whole-op fallback。",
35+ },
36+ {
37+ "path": "mindspeed_ops/arch32/triton/mamba3/mamba3_siso_bwd_impl.py",
38+ "kind": "production kernel",
39+ "purpose": "Triton-Ascend kernel 和 host dispatch；包含 rotary_bias_angles、dk_state_post、ddt_dtrap_dinput_states 三个 kernel。",
40+ "review_check": "生产计算必须进入本文件内 Triton kernel；dispatch 只能按运行时 shape/dtype/metadata。",
41+ },
42+ {
43+ "path": "mindspeed_ops/arch32/triton/mamba3/mamba3_utils.py",
44+ "kind": "production helper",
45+ "purpose": "Triton JIT helper：rotary 小角度 sin/cos 多项式、sigmoid、silu。",
46+ "review_check": "仅提供 JIT helper，不读取外部 reference、case、golden 或上游包。",
47+ },
48+ {
49+ "path": "tests/unit_tests/triton/test_mamba3_siso_bwd.py",
50+ "kind": "unit test",
51+ "purpose": "pytest 精度单测，覆盖 compute_dqktheta 和 compute_ddt_dtrap_dinput_states 两条生产路径。",
52+ "review_check": "测试侧可调用 torch reference；生产入口不能反向依赖测试代码。",
53+ },
54+ {
55+ "path": "tests/atk_tests/triton/mamba3_siso_bwd/generate_mamba3_siso_bwd.py",
56+ "kind": "ATK generator",
57+ "purpose": "ATK 用例生成器；从提交内参数空间生成 shape/dtype/attr。",
58+ "review_check": "不读取外部 task cases、answer key 或私有 workload 文件。",
59+ },
60+ {
61+ "path": "tests/atk_tests/triton/mamba3_siso_bwd/mamba3_siso_bwd.yaml",
62+ "kind": "ATK config",
63+ "purpose": "ATK 算子配置，声明输入、dtype、取值范围、baseline API 与 Triton API。",
64+ "review_check": "配置中的有效范围要与本交付支持边界一致。",
65+ },
66+ {
67+ "path": "tests/atk_tests/triton/mamba3_siso_bwd/triton_mamba3_siso_bwd.py",
68+ "kind": "ATK wrapper",
69+ "purpose": "ATK baseline/candidate wrapper；baseline 仅用于测试侧对比。",
70+ "review_check": "candidate wrapper 必须调用 `mindspeed_ops.api.triton.mamba3_siso_bwd`。",
71+ },
72+ {
73+ "path": "tests/atk_tests/triton/mamba3_siso_bwd/reference_impl.py",
74+ "kind": "test reference",
75+ "purpose": "提交内 torch reference，用于 UT/ATK/自证脚本精度校验。",
76+ "review_check": "只能被测试和自证脚本引用，生产路径不能 import 它。",
77+ },
78+ {
79+ "path": "tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py",
80+ "kind": "self-check script",
81+ "purpose": "一键生成 Markdown/JSON 自证报告，包含文件范围、静态 no-fallback、live 精度。",
82+ "review_check": "只依赖 PR 内文件和当前 MindSpeed-Ops 运行环境；不读取外部评测目录。",
83+ },
84+ {
85+ "path": "docs/triton/mamba3_siso_bwd.md",
86+ "kind": "documentation",
87+ "purpose": "中文实践说明和审查说明，记录支持边界、验证命令和证据口径。",
88+ "review_check": "文档声明必须能由本脚本或仓内测试复核。",
89+ },
90+]
91+ 
92+PRODUCTION_FILES = [item["path"] for item in FILE_INVENTORY if item["kind"].startswith("production")]
93+ 
94+REVIEW_CASES_DQKTHETA = [
95+ {"case_id": 1, "batch": 1, "seqlen": 32, "nheads_qk": 1, "nheads": 1, "headdim_qk": 64, "dtype": "float16", "chunk_size": 64},
96+ {"case_id": 2, "batch": 1, "seqlen": 64, "nheads_qk": 1, "nheads": 2, "headdim_qk": 64, "dtype": "float16", "chunk_size": 64},
97+ {"case_id": 3, "batch": 2, "seqlen": 64, "nheads_qk": 2, "nheads": 4, "headdim_qk": 64, "dtype": "float16", "chunk_size": 32},
98+ {"case_id": 5, "batch": 1, "seqlen": 96, "nheads_qk": 1, "nheads": 2, "headdim_qk": 64, "dtype": "bfloat16", "chunk_size": 64},
99+ {"case_id": 6, "batch": 2, "seqlen": 96, "nheads_qk": 1, "nheads": 2, "headdim_qk": 64, "dtype": "bfloat16", "chunk_size": 64},
100+ {"case_id": 7, "batch": 1, "seqlen": 128, "nheads_qk": 2, "nheads": 4, "headdim_qk": 64, "dtype": "float16", "chunk_size": 64},
101+]
102+ 
103+REVIEW_CASES_DDT = [
104+ {"case_id": 101, "batch": 1, "nheads": 1, "seqlen": 32, "dtype": "float16"},
105+ {"case_id": 102, "batch": 1, "nheads": 2, "seqlen": 64, "dtype": "float16"},
106+ {"case_id": 103, "batch": 2, "nheads": 2, "seqlen": 64, "dtype": "float16"},
107+ {"case_id": 104, "batch": 1, "nheads": 4, "seqlen": 96, "dtype": "bfloat16"},
108+ {"case_id": 105, "batch": 2, "nheads": 4, "seqlen": 128, "dtype": "float16"},
109+]
110+ 
111+QUICK_CASE_IDS = {1, 101}
112+ABS_TOL = 1e-3
113+REL_RMSE_TOL = 1e-2
114+ 
115+ 
116+def _round(value: float, digits: int = 6) -> float:
117+ return round(float(value), digits)
118+ 
119+ 
120+def _max_abs(actual: torch.Tensor, expected: torch.Tensor) -> float:
121+ return (actual.detach().cpu().float() - expected.detach().cpu().float()).abs().max().item()
122+ 
123+ 
124+def _relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float:
125+ diff = actual.detach().cpu().float().flatten() - expected.detach().cpu().float().flatten()
126+ base = expected.detach().cpu().float().flatten()
127+ return torch.sqrt(torch.mean(diff * diff)).item() / (torch.sqrt(torch.mean(base * base)).item() + 1e-8)
128+ 
129+ 
130+def _dtype(name: str) -> torch.dtype:
131+ return {"float16": torch.float16, "bfloat16": torch.bfloat16}[name]
132+ 
133+ 
134+def _has_npu() -> bool:
135+ return hasattr(torch, "npu") and torch.npu.is_available()
136+ 
137+ 
138+def _sync() -> None:
139+ if _has_npu():
140+ torch.npu.synchronize()
141+ 
142+ 
143+def _scan_imports(path: Path) -> list[str]:
144+ tree = ast.parse(path.read_text(encoding="utf-8"))
145+ imports: list[str] = []
146+ for node in ast.walk(tree):
147+ if isinstance(node, ast.Import):
148+ imports.extend(alias.name for alias in node.names)
149+ elif isinstance(node, ast.ImportFrom) and node.module:
150+ imports.append(node.module)
151+ return imports
152+ 
153+ 
154+def static_no_fallback_check() -> dict[str, Any]:
155+ missing_files = [item["path"] for item in FILE_INVENTORY if not (REPO_ROOT / item["path"]).is_file()]
156+ forbidden_import_fragments = [
157+ "state_spaces", "state-spaces", "mamba_ssm", "golden", "reference",
158+ "tests.unit_tests", "tests.atk_tests",
159+ ]
160+ forbidden_text_fragments = [
161+ "torch.ops", "torch_npu.ops", "torch_npu.contrib", "torch_npu.npu_",
162+ "mamba_ssm", "task golden files", "refer/upstreams", "state-spaces-mamba",
163+ "--external-evidence-root",
164+ ]
165+ allowed_runtime_fragments = {
166+ "torch": "tensor/runtime API for input metadata and device checks",
167+ "torch_npu": "runtime availability probe only; no torch_npu whole-op call is allowed",
168+ "triton": "Triton-Ascend JIT and launch API",
169+ "triton.language": "Triton language primitives",
170+ }
171+ file_results = []
172+ violations = [f"missing file: {item}" for item in missing_files]
173+ for rel in PRODUCTION_FILES:
174+ path = REPO_ROOT / rel
175+ if not path.is_file():
176+ continue
177+ text = path.read_text(encoding="utf-8")
178+ imports = _scan_imports(path)
179+ file_violations = []
180+ for item in imports:
181+ if any(fragment in item for fragment in forbidden_import_fragments):
182+ file_violations.append(f"forbidden import: {item}")
183+ for fragment in forbidden_text_fragments:
184+ if fragment in text:
185+ file_violations.append(f"forbidden text: {fragment}")
186+ file_results.append({"path": rel, "imports": imports, "violations": file_violations})
187+ violations.extend(f"{rel}: {v}" for v in file_violations)
188+ 
189+ return {
190+ "passed": not violations,
191+ "file_inventory": FILE_INVENTORY,
192+ "missing_files": missing_files,
193+ "production_files": PRODUCTION_FILES,
194+ "allowed_runtime_fragments": allowed_runtime_fragments,
195+ "file_results": file_results,
196+ "violations": violations,
197+ }
198+ 
199+ 
200+def _run_dqktheta_case(case: dict[str, Any], device: torch.device):
201+ inputs = make_inputs_dqktheta(
202+ case["batch"], case["seqlen"], case["nheads_qk"], case["nheads"],
203+ case["headdim_qk"], _dtype(case["dtype"]), device,
204+ )
205+ actual = compute_dqktheta(*inputs, chunk_size=case["chunk_size"])
206+ expected = ref_compute_dqktheta(*inputs, chunk_size=case["chunk_size"])
207+ _sync()
208+ 
209+ name = f"dqktheta_c{case['case_id']}"
210+ results = []
211+ for i, (tensor_name, act, exp) in enumerate(zip(
212+ ["dq", "dk", "dangles", "dscale", "dgamma"],
213+ actual[:5], expected[:5],
214+ )):
215+ max_abs = _max_abs(act, exp)
216+ relative_rmse = _relative_rmse(act, exp)
217+ mare = get_mare(act, exp)
218+ mere = get_mere(act, exp)
219+ rmse = get_rmse(act, exp)
220+ accuracy_passed = max_abs <= ABS_TOL or relative_rmse <= REL_RMSE_TOL
221+ results.append({
222+ "tensor": tensor_name, "max_abs": _round(max_abs, 10),
223+ "relative_rmse": _round(relative_rmse, 10),
224+ "mare": _round(mare, 10), "mere": _round(mere, 10),
225+ "rmse": _round(rmse, 10), "accuracy_passed": accuracy_passed,
226+ })
227+ 
228+ all_passed = all(r["accuracy_passed"] for r in results)
229+ return {
230+ "case_id": case["case_id"], "dtype": case["dtype"],
231+ "shape": f"B{case['batch']}_S{case['seqlen']}_Hq{case['nheads_qk']}_H{case['nheads']}_K{case['headdim_qk']}_CS{case['chunk_size']}",
232+ "results": results, "passed": all_passed,
233+ }
234+ 
235+ 
236+def _run_ddt_case(case: dict[str, Any], device: torch.device):
237+ inputs = make_inputs_ddt(
238+ case["batch"], case["nheads"], case["seqlen"], _dtype(case["dtype"]), device,
239+ )
240+ actual_ddt, actual_dtrap, *_ = compute_ddt_dtrap_dinput_states(*inputs)
241+ expected_ddt, expected_dtrap = ref_compute_ddt_dtrap_dinput_states(*inputs)
242+ _sync()
243+ 
244+ results = []
245+ for tensor_name, act, exp in [("dDT", actual_ddt, expected_ddt), ("dTrap", actual_dtrap, expected_dtrap)]:
246+ max_abs = _max_abs(act, exp)
247+ relative_rmse = _relative_rmse(act, exp)
248+ mare = get_mare(act, exp)
249+ mere = get_mere(act, exp)
250+ rmse = get_rmse(act, exp)
251+ accuracy_passed = max_abs <= ABS_TOL or relative_rmse <= REL_RMSE_TOL
252+ results.append({
253+ "tensor": tensor_name, "max_abs": _round(max_abs, 10),
254+ "relative_rmse": _round(relative_rmse, 10),
255+ "mare": _round(mare, 10), "mere": _round(mare, 10),
256+ "rmse": _round(rmse, 10), "accuracy_passed": accuracy_passed,
257+ })
258+ 
259+ all_passed = all(r["accuracy_passed"] for r in results)
260+ return {
261+ "case_id": case["case_id"], "dtype": case["dtype"],
262+ "shape": f"B{case['batch']}_H{case['nheads']}_S{case['seqlen']}",
263+ "results": results, "passed": all_passed,
264+ }
265+ 
266+ 
267+def live_case_check(cases_dqktheta: list, cases_ddt: list) -> dict[str, Any]:
268+ if not _has_npu():
269+ return {"passed": False, "skipped": True, "reason": "NPU is not available", "cases": [], "metrics": None}
270+ 
271+ device = torch.device("npu")
272+ all_rows = []
273+ 
274+ for case in cases_dqktheta:
275+ row = _run_dqktheta_case(case, device)
276+ all_rows.append(row)
277+ 
278+ for case in cases_ddt:
279+ row = _run_ddt_case(case, device)
280+ all_rows.append(row)
281+ 
282+ passed = all(row["passed"] for row in all_rows)
283+ return {"passed": passed, "skipped": False, "cases": all_rows, "metrics": None}
284+ 
285+ 
286+def build_report(quick: bool) -> dict[str, Any]:
287+ static = static_no_fallback_check()
288+ dqktheta_cases = [case for case in REVIEW_CASES_DQKTHETA if not quick or case["case_id"] in QUICK_CASE_IDS]
289+ ddt_cases = [case for case in REVIEW_CASES_DDT if not quick or case["case_id"] in QUICK_CASE_IDS]
290+ live = live_case_check(dqktheta_cases, ddt_cases)
291+ passed = static["passed"] and live["passed"]
292+ return {
293+ "passed": passed,
294+ "mode": "quick" if quick else "full",
295+ "accuracy_rule": {
296+ "abs_tol": ABS_TOL,
297+ "relative_rmse_tol": REL_RMSE_TOL,
298+ "pass_rule": "max_abs <= abs_tol OR relative_rmse <= relative_rmse_tol",
299+ },
300+ "static_no_fallback": static,
301+ "live_cases": live,
302+ "known_boundary": {
303+ "arch": "arch32 production path; arch35 fails loudly",
304+ "operator_scope": "dense non-varlen backward path (dqktheta, ddt_dtrap)",
305+ "unsupported_scope": ["varlen", "state-return", "full backward combined"],
306+ "known_limitation": "headdim_qk=128 triggers NPU UB overflow (loop unrolling doubles buffer usage); current PR covers headdim_qk=64 only",
307+ },
308+ }
309+ 
310+ 
311+def render_markdown(report: dict[str, Any]) -> str:
312+ lines = [
313+ "# mamba3_siso_bwd 自证报告",
314+ "",
315+ "## 0. 运行与结论",
316+ "",
317+ f"- overall: {'PASS' if report['passed'] else 'FAIL'}",
318+ f"- mode: `{report['mode']}`",
319+ f"- static_no_fallback: {'PASS' if report['static_no_fallback']['passed'] else 'FAIL'}",
320+ f"- live_cases: {'PASS' if report['live_cases']['passed'] else 'FAIL'}",
321+ "- repo_root: `<MindSpeed-Ops repository root>`",
322+ "- external_workspace_required: `no`",
323+ "",
324+ "## 1. 提交范围",
325+ "",
326+ "| 文件 | 类型 | 用途 | 审查项 | 存在 |",
327+ "| --- | --- | --- | --- | --- |",
328+ ]
329+ for item in report["static_no_fallback"]["file_inventory"]:
330+ path = item["path"]
331+ exists = (REPO_ROOT / path).is_file()
332+ lines.append(
333+ f"| `{path}` | {item['kind']} | {item['purpose']} | {item['review_check']} | "
334+ f"{'yes' if exists else 'NO'} |"
335+ )
336+ 
337+ lines.extend(["", "## 2. 无 fallback 静态检查", ""])
338+ lines.extend(["允许的生产依赖边界：", "", "| import/text | 用途 |", "| --- | --- |"])
339+ for name, reason in report["static_no_fallback"]["allowed_runtime_fragments"].items():
340+ lines.append(f"| `{name}` | {reason} |")
341+ lines.extend(["", "生产文件 import 扫描：", ""])
342+ for item in report["static_no_fallback"]["file_results"]:
343+ imports = ", ".join(f"`{name}`" for name in item["imports"]) or "(none)"
344+ lines.append(f"- `{item['path']}` imports: {imports}")
345+ lines.append("")
346+ if report["static_no_fallback"]["violations"]:
347+ for item in report["static_no_fallback"]["violations"]:
348+ lines.append(f"- FAIL: {item}")
349+ else:
350+ lines.append("- PASS: production files do not import or reference torch_npu high-level ops, golden/reference code, upstream runtime packages, peer code, or other backend fallback.")
351+ 
352+ lines.extend(["", "## 3. live 精度", ""])
353+ live = report["live_cases"]
354+ if live.get("skipped"):
355+ lines.append(f"- SKIPPED: {live['reason']}")
356+ else:
357+ lines.append("| case | dtype | shape | tensor | max_abs | rel_rmse | MARE | MERE | RMSE | result |")
358+ lines.append("| ---: | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | --- |")
359+ for item in live["cases"]:
360+ for r in item["results"]:
361+ lines.append(
362+ f"| {item['case_id']} | {item['dtype']} | `{item['shape']}` | {r['tensor']} "
363+ f"| {r['max_abs']:.10f} | {r['relative_rmse']:.10f} "
364+ f"| {r['mare']:.10f} | {r['mere']:.10f} | {r['rmse']:.10f} "
365+ f"| {'PASS' if r['accuracy_passed'] else 'FAIL'} |"
366+ )
367+ 
368+ boundary = report["known_boundary"]
369+ lines.extend([
370+ "",
371+ "## 4. 已知边界",
372+ "",
373+ f"- arch: {boundary['arch']}",
374+ f"- operator_scope: {boundary['operator_scope']}",
375+ "- unsupported_scope: " + ", ".join(boundary["unsupported_scope"]),
376+ f"- known_limitation: {boundary.get('known_limitation', '(none)')}",
377+ ])
378+ return "\n".join(lines) + "\n"
379+ 
380+ 
381+def parse_args() -> argparse.Namespace:
382+ parser = argparse.ArgumentParser(description="Generate self-contained mamba3_siso_bwd review report.")
383+ parser.add_argument("--quick", action="store_true", help="Run only two representative cases.")
384+ parser.add_argument("--output", type=Path, default=None, help="Optional markdown output path.")
385+ parser.add_argument("--json", action="store_true", help="Print JSON instead of markdown.")
386+ return parser.parse_args()
387+ 
388+ 
389+def main() -> int:
390+ args = parse_args()
391+ report = build_report(quick=args.quick)
392+ text = json.dumps(report, indent=2, ensure_ascii=False) if args.json else render_markdown(report)
393+ if args.output:
394+ args.output.write_text(text, encoding="utf-8")
395+ else:
396+ print(text, end="")
397+ return 0 if report["passed"] else 1
398+ 
399+ 
400+if __name__ == "__main__":
401+ raise SystemExit(main())
Atests/atk_tests/triton/mamba3_siso_bwd/triton_mamba3_siso_bwd.py+75-0
@@ -0,0 +1,75 @@
1+import torch
2+ 
3+from atk.configs.dataset_config import InputDataset
4+from atk.tasks.api_execute import register
5+from atk.tasks.api_execute.base_api import BaseApi
6+from atk.tasks.api_execute.triton_base_api import TritonBaseApi
7+from mindspeed_ops.api.triton.mamba3_siso_bwd import compute_dqktheta, compute_ddt_dtrap_dinput_states
8+from tests.atk_tests.triton.mamba3_siso_bwd.reference_impl import (
9+ ref_compute_dqktheta,
10+ ref_compute_ddt_dtrap_dinput_states,
11+)
12+ 
13+ 
14+def _get_dqktheta_inputs(input_data: InputDataset):
15+ kwargs = input_data.kwargs
16+ q = kwargs.get("q")
17+ k = kwargs.get("k")
18+ scale = kwargs.get("scale")
19+ gamma = kwargs.get("gamma")
20+ q_bias = kwargs.get("q_bias")
21+ k_bias = kwargs.get("k_bias")
22+ angles = kwargs.get("angles")
23+ dq_in = kwargs.get("dq_in")
24+ dk_in = kwargs.get("dk_in")
25+ dqk = kwargs.get("dqk")
26+ chunk_size = kwargs.get("chunk_size", 64)
27+ return q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, chunk_size
28+ 
29+ 
30+def _get_ddt_inputs(input_data: InputDataset):
31+ kwargs = input_data.kwargs
32+ dscale = kwargs.get("dscale")
33+ dgamma = kwargs.get("dgamma")
34+ dt = kwargs.get("dt")
35+ trap = kwargs.get("trap")
36+ return dscale, dgamma, dt, trap
37+ 
38+ 
39+@register("torch_mamba3_siso_bwd_dqktheta")
40+class TorchMamba3SisoBwdDqkthetaApi(BaseApi):
41+ def __call__(self, input_data: InputDataset, with_output: bool = False):
42+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, chunk_size = _get_dqktheta_inputs(input_data)
43+ return ref_compute_dqktheta(
44+ q=q, k=k, scale=scale, gamma=gamma, q_bias=q_bias, k_bias=k_bias,
45+ angles=angles, dq_in=dq_in, dk_in=dk_in, dqk=dqk, chunk_size=chunk_size,
46+ )
47+ 
48+ 
49+@register("triton_mamba3_siso_bwd_dqktheta")
50+class TritonMamba3SisoBwdDqkthetaApi(TritonBaseApi):
51+ def __call__(self, input_data: InputDataset, with_output: bool = False):
52+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, chunk_size = _get_dqktheta_inputs(input_data)
53+ return compute_dqktheta(
54+ q=q, k=k, scale=scale, gamma=gamma, q_bias=q_bias, k_bias=k_bias,
55+ angles=angles, dq_in=dq_in, dk_in=dk_in, dqk=dqk, chunk_size=chunk_size,
56+ )
57+ 
58+ 
59+@register("torch_mamba3_siso_bwd_ddt")
60+class TorchMamba3SisoBwdDdtApi(BaseApi):
61+ def __call__(self, input_data: InputDataset, with_output: bool = False):
62+ dscale, dgamma, dt, trap = _get_ddt_inputs(input_data)
63+ return ref_compute_ddt_dtrap_dinput_states(
64+ dscale=dscale, dgamma=dgamma, dt=dt, trap=trap,
65+ )
66+ 
67+ 
68+@register("triton_mamba3_siso_bwd_ddt")
69+class TritonMamba3SisoBwdDdtApi(TritonBaseApi):
70+ def __call__(self, input_data: InputDataset, with_output: bool = False):
71+ dscale, dgamma, dt, trap = _get_ddt_inputs(input_data)
72+ ddt, dtrap, *_ = compute_ddt_dtrap_dinput_states(
73+ dscale=dscale, dgamma=dgamma, dt=dt, trap=trap,
74+ )
75+ return ddt, dtrap
Atests/unit_tests/triton/test_mamba3_siso_bwd.py+262-0
@@ -0,0 +1,262 @@
1+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
2+# Copyright (c) 2025, Dao AI Lab, Goombalab
3+ 
4+import pytest
5+import torch
6+ 
7+from mindspeed_ops.api.triton.mamba3_siso_bwd import compute_dqktheta, compute_ddt_dtrap_dinput_states
8+from tests.utils import assert_close
9+ 
10+ 
11+# =============================================================================
12+# Reference implementations
13+# =============================================================================
14+ 
15+def _apply_rotary(x: torch.Tensor, angles: torch.Tensor) -> torch.Tensor:
16+ """Apply rotary embedding to x using angles."""
17+ out = x.to(torch.float32).clone()
18+ rotary_dim = min(angles.shape[-1], out.shape[-1] // 2)
19+ if rotary_dim <= 0:
20+ return out
21+ cos = torch.cos(angles[..., :rotary_dim].to(torch.float32))
22+ sin = torch.sin(angles[..., :rotary_dim].to(torch.float32))
23+ even = out[..., 0:2 * rotary_dim:2].clone()
24+ odd = out[..., 1:2 * rotary_dim:2].clone()
25+ out[..., 0:2 * rotary_dim:2] = even * cos - odd * sin
26+ out[..., 1:2 * rotary_dim:2] = even * sin + odd * cos
27+ return out
28+ 
29+ 
30+def _inv_apply_rotary(x: torch.Tensor, angles: torch.Tensor) -> torch.Tensor:
31+ """Apply inverse rotary embedding to x using angles."""
32+ out = x.to(torch.float32).clone()
33+ rotary_dim = min(angles.shape[-1], out.shape[-1] // 2)
34+ if rotary_dim <= 0:
35+ return out
36+ cos = torch.cos(angles[..., :rotary_dim].to(torch.float32))
37+ sin = torch.sin(angles[..., :rotary_dim].to(torch.float32))
38+ even = out[..., 0:2 * rotary_dim:2].clone()
39+ odd = out[..., 1:2 * rotary_dim:2].clone()
40+ out[..., 0:2 * rotary_dim:2] = even * cos + odd * sin
41+ out[..., 1:2 * rotary_dim:2] = -even * sin + odd * cos
42+ return out
43+ 
44+ 
45+def ref_compute_dqktheta(
46+ q: torch.Tensor,
47+ k: torch.Tensor,
48+ scale: torch.Tensor,
49+ gamma: torch.Tensor,
50+ q_bias: torch.Tensor,
51+ k_bias: torch.Tensor,
52+ angles: torch.Tensor,
53+ dq_in: torch.Tensor,
54+ dk_in: torch.Tensor,
55+ dqk: torch.Tensor,
56+ chunk_size: int = 64,
57+) -> tuple:
58+ """Reference implementation of compute_dqktheta."""
59+ batch, seqlen, nheads_qk, headdim_qk = q.shape
60+ nheads = scale.shape[1]
61+ GQA_RATIO = nheads // nheads_qk
62+ work_dtype = torch.float32
63+ 
64+ dq = torch.zeros((batch, seqlen, nheads_qk, headdim_qk), device=q.device, dtype=work_dtype)
65+ dk = torch.zeros((batch, seqlen, nheads_qk, headdim_qk), device=k.device, dtype=work_dtype)
66+ dangles = torch.zeros((batch, seqlen, nheads, angles.shape[-1]), device=angles.device, dtype=work_dtype)
67+ dscale = torch.zeros((batch, nheads, seqlen), device=scale.device, dtype=work_dtype)
68+ dgamma = torch.zeros((batch, nheads, seqlen), device=gamma.device, dtype=work_dtype)
69+ dq_bias = torch.zeros((nheads, headdim_qk), device=q.device, dtype=work_dtype)
70+ dk_bias = torch.zeros((nheads, headdim_qk), device=k.device, dtype=work_dtype)
71+ 
72+ for b_idx in range(batch):
73+ for qk_head_idx in range(nheads_qk):
74+ for gqa_idx in range(GQA_RATIO):
75+ nhead_idx = qk_head_idx * GQA_RATIO + gqa_idx
76+ 
77+ q_pre = q[b_idx, :, qk_head_idx, :].to(work_dtype) + q_bias[nhead_idx].to(work_dtype)
78+ k_pre = k[b_idx, :, qk_head_idx, :].to(work_dtype) + k_bias[nhead_idx].to(work_dtype)
79+ dq_in_h = dq_in[b_idx, :, nhead_idx, :].to(work_dtype)
80+ dk_in_h = dk_in[b_idx, :, nhead_idx, :].to(work_dtype)
81+ dqk_h = dqk[b_idx, nhead_idx, :].to(work_dtype)
82+ scale_h = scale[b_idx, nhead_idx, :].to(work_dtype)
83+ gamma_h = gamma[b_idx, nhead_idx, :].to(work_dtype)
84+ angles_h = angles[b_idx, :, nhead_idx, :].to(work_dtype)
85+ 
86+ # dGamma = dQK * (Q_wbias . K_wbias)
87+ qk_dot = (q_pre * k_pre).sum(dim=-1)
88+ dgamma[b_idx, nhead_idx, :] = dqk_h * qk_dot
89+ 
90+ # cos/sin
91+ cos_ang = torch.cos(angles_h)
92+ sin_ang = torch.sin(angles_h)
93+ 
94+ # dScale = sum(dK_in * K_rot)
95+ k_rot = _apply_rotary(k_pre, angles_h)
96+ dscale[b_idx, nhead_idx, :] = (dk_in_h * k_rot).sum(dim=-1)
97+ 
98+ # Inverse rotary
99+ dk_in_scaled = dk_in_h * scale_h.unsqueeze(-1)
100+ dq_unrot = _inv_apply_rotary(dq_in_h, angles_h)
101+ dk_unrot = _inv_apply_rotary(dk_in_scaled, angles_h)
102+ 
103+ # Add dQK path
104+ dqk_scaled = (dqk_h * gamma_h).unsqueeze(-1)
105+ dq_pre = dq_unrot + dqk_scaled * k_pre
106+ dk_pre = dk_unrot + dqk_scaled * q_pre
107+ 
108+ # Accumulate GQA
109+ dq[b_idx, :, qk_head_idx, :] += dq_pre
110+ dk[b_idx, :, qk_head_idx, :] += dk_pre
111+ dq_bias[nhead_idx, :] += dq_pre.sum(dim=0)
112+ dk_bias[nhead_idx, :] += dk_pre.sum(dim=0)
113+ 
114+ # dAngles
115+ rotary_dim = min(angles_h.shape[-1], headdim_qk // 2)
116+ if rotary_dim > 0:
117+ even_q = q_pre[:, 0:2 * rotary_dim:2]
118+ odd_q = q_pre[:, 1:2 * rotary_dim:2]
119+ even_k = k_pre[:, 0:2 * rotary_dim:2]
120+ odd_k = k_pre[:, 1:2 * rotary_dim:2]
121+ even_dq = dq_in_h[:, 0:2 * rotary_dim:2]
122+ odd_dq = dq_in_h[:, 1:2 * rotary_dim:2]
123+ even_dk = dk_in_scaled[:, 0:2 * rotary_dim:2]
124+ odd_dk = dk_in_scaled[:, 1:2 * rotary_dim:2]
125+ 
126+ dtheta_q = even_dq * (-even_q * sin_ang - odd_q * cos_ang) + \
127+ odd_dq * (even_q * cos_ang - odd_q * sin_ang)
128+ dtheta_k = even_dk * (-even_k * sin_ang - odd_k * cos_ang) + \
129+ odd_dk * (even_k * cos_ang - odd_k * sin_ang)
130+ dangles[b_idx, :, nhead_idx, :rotary_dim] = dtheta_q + dtheta_k
131+ 
132+ return dq, dk, dq_bias, dk_bias, dangles, dscale, dgamma
133+ 
134+ 
135+def ref_compute_ddt_dtrap_dinput_states(
136+ dscale: torch.Tensor,
137+ dgamma: torch.Tensor,
138+ dt: torch.Tensor,
139+ trap: torch.Tensor,
140+) -> tuple:
141+ """Reference implementation of compute_ddt_dtrap_dinput_states."""
142+ batch, nheads, seqlen = dscale.shape
143+ work_dtype = torch.float32
144+ 
145+ dDT = torch.zeros_like(dt, dtype=work_dtype)
146+ dTrap = torch.zeros_like(trap, dtype=work_dtype)
147+ 
148+ for b_idx in range(batch):
149+ for h_idx in range(nheads):
150+ dscale_h = dscale[b_idx, h_idx].to(work_dtype)
151+ dgamma_h = dgamma[b_idx, h_idx].to(work_dtype)
152+ dt_h = dt[b_idx, h_idx].to(work_dtype)
153+ trap_h = torch.sigmoid(trap[b_idx, h_idx].to(work_dtype))
154+ 
155+ for t in range(seqlen):
156+ dscale_prev = dscale_h[t - 1] if t > 0 else torch.tensor(0.0, device=dt.device, dtype=work_dtype)
157+ ddt_t = (dgamma_h[t] + dscale_h[t]) * trap_h[t] + dscale_prev * (1.0 - trap_h[t])
158+ dtrap_t = (dgamma_h[t] + dscale_h[t]) * dt_h[t] - dscale_prev * dt_h[t]
159+ dtrap_presig_t = dtrap_t * trap_h[t] * (1.0 - trap_h[t])
160+ 
161+ dDT[b_idx, h_idx, t] = ddt_t
162+ dTrap[b_idx, h_idx, t] = dtrap_presig_t
163+ 
164+ return dDT, dTrap
165+ 
166+ 
167+def make_inputs(batch, seqlen, nheads_qk, nheads, headdim_qk, dtype, device):
168+ """Create test inputs for the backward operators."""
169+ torch.manual_seed(2026 + batch + seqlen + nheads + headdim_qk)
170+ headdim_angles = headdim_qk // 2
171+ 
172+ q = (torch.rand(batch, seqlen, nheads_qk, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
173+ k = (torch.rand(batch, seqlen, nheads_qk, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
174+ scale = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.5 + 0.5
175+ gamma = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.3
176+ q_bias = (torch.rand(nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.1
177+ k_bias = (torch.rand(nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.1
178+ angles = torch.rand(batch, seqlen, nheads, headdim_angles, device=device, dtype=torch.float32) * 0.25
179+ dq_in = (torch.rand(batch, seqlen, nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
180+ dk_in = (torch.rand(batch, seqlen, nheads, headdim_qk, device=device, dtype=dtype) - 0.5) * 0.4
181+ dqk = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.2
182+ dt = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 0.1
183+ trap = torch.rand(batch, nheads, seqlen, device=device, dtype=torch.float32) * 2.0 - 1.0
184+ 
185+ return q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, dt, trap
186+ 
187+ 
188+# =============================================================================
189+# Test cases
190+# =============================================================================
191+ 
192+TEST_CASES_DQKTHETA = [
193+ pytest.param(1, 32, 1, 1, 64, torch.float16, 64, id="fp16_small"),
194+ pytest.param(1, 64, 1, 2, 64, torch.float16, 64, id="fp16_gqa"),
195+ pytest.param(1, 96, 1, 2, 64, torch.bfloat16, 64, id="bf16_multi_chunk"),
196+ pytest.param(2, 64, 2, 4, 64, torch.float16, 32, id="fp16_batch2"),
197+]
198+ 
199+TEST_CASES_DDT = [
200+ pytest.param(1, 32, 1, 1, 64, torch.float16, 64, id="ddt_small"),
201+ pytest.param(1, 64, 1, 2, 64, torch.float16, 64, id="ddt_gqa"),
202+ pytest.param(2, 96, 1, 2, 64, torch.bfloat16, 64, id="ddt_bf16_batch2"),
203+]
204+ 
205+ 
206+@pytest.mark.parametrize(
207+ "batch,seqlen,nheads_qk,nheads,headdim_qk,dtype,chunk_size",
208+ TEST_CASES_DQKTHETA,
209+)
210+def test_compute_dqktheta(batch, seqlen, nheads_qk, nheads, headdim_qk, dtype, chunk_size):
211+ if not hasattr(torch, "npu") or not torch.npu.is_available():
212+ pytest.skip("NPU is not available")
213+ device = torch.device("npu")
214+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, dt, trap = make_inputs(
215+ batch, seqlen, nheads_qk, nheads, headdim_qk, dtype, device
216+ )
217+ 
218+ actual = compute_dqktheta(
219+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk,
220+ chunk_size=chunk_size,
221+ )
222+ 
223+ expected = ref_compute_dqktheta(
224+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, chunk_size=chunk_size,
225+ )
226+ 
227+ for i, (name, act, exp) in enumerate(zip(
228+ ["dq", "dk", "dangles", "dscale", "dgamma"],
229+ actual[:5],
230+ expected[:5],
231+ )):
232+ assert_close(f"compute_dqktheta_{name}", exp, act, 1e-3, err_atol=1e-3)
atomgit-botatomgit-bot
atomgit-botatomgit-bot7月12日

🟡 Medium Priority

changed line: tests/unit_tests/triton/test_mamba3_siso_bwd.py 第 227-232 行，test_compute_dqktheta 只校验了 actual[:5]（dq, dk, dangles, dscale, dgamma），跳过了 dq_bias（索引 2）和 dk_bias（索引 3）。

affected behavior: compute_dqktheta 返回 7 个 tensor，但单元测试和 self_check 脚本均未验证 dq_bias 和 dk_bias 的精度。这两个 tensor 由 kernel 内 even/odd 累加并 partial sum 输出（impl 第 620-621 行），其计算路径涉及 bias 加载、GQA 循环内的逐 head 累加和 store，若 kernel 中 bias 索引或 stride 有误，现有测试无法发现。

failure mode: 若 kernel 内 dq_bias/dk_bias 计算有误（如 stride 错误、GQA 分组索引偏移），用户在生产环境中得到的 bias 梯度将是错误的，但所有测试仍会 PASS。

suggested fix: 在 test_compute_dqktheta 中添加对 actual[2]（dq_bias）和 actual[3]（dk_bias）的 assert_close 校验；在 self_check_mamba3_siso_bwd.py 的 _run_dqktheta_case 中同样加入 dq_bias 和 dk_bias 的精度对比。

建议：在 test_compute_dqktheta 和 self_check 的 _run_dqktheta_case 中补齐 dq_bias 和 dk_bias 的精度校验

likedislike
不准确？
atomgit-botatomgit-bot7月12日

🟡 Medium Priority

changed line: tests/unit_tests/triton/test_mamba3_siso_bwd.py 第 227-232 行和 tests/atk_tests/triton/mamba3_siso_bwd/self_check_mamba3_siso_bwd.py 第 211-213 行。

compute_dqktheta 返回 7 个 tensor：(dq, dk, dq_bias, dk_bias, dangles, dscale, dgamma)。但测试代码使用 actual[:5] 取前 5 个，与名称列表 ["dq", "dk", "dangles", "dscale", "dgamma"] 做 zip，导致：

  1. 名称标签错位：索引 2 的 dq_bias 被标记为 "dangles"，索引 3 的 dk_bias 被标记为 "dscale"，索引 4 的 dangles 被标记为 "dgamma"
  2. dscale 和 dgamma 完全未测试：这两个核心输出（索引 5、6）不在 actual[:5] 范围内，从未参与精度校验
  3. 自证报告误导：docs/triton/mamba3_siso_bwd.md 第 155-200 行展示的 "dscale" 和 "dgamma" 精度数据，实际来自 dk_bias 和 dangles 的比较结果，而非真正的 dscale/dgamma

affected behavior: dscale 和 dgamma 是 kernel 的关键输出，它们直接作为 compute_ddt_dtrap_dinput_states 的输入。若 kernel 中 dscale/dgamma 计算有误（如 GQA reduction 错误、stride 偏移），所有测试仍会 PASS，但下游 dDT/dTrap 计算将基于错误的 dscale/dgamma。

failure mode: dscale/dgamma 的错误会传播到 ddt/dtrap 计算，但测试无法发现。自证报告中的 "dscale"/"dgamma" 行实际展示的是 dk_bias/dangles 精度，造成虚假信心。

suggested fix: 修正为 actual 与 expected 的完整 7 元组对比，名称列表改为 ["dq", "dk", "dq_bias", "dk_bias", "dangles", "dscale", "dgamma"]。

建议：修正名称列表与实际返回值对齐，补齐 dscale 和 dgamma 的精度校验

改动建议
232
+ for i, (name, act, exp) in enumerate(zip(
233
+ ["dq", "dk", "dq_bias", "dk_bias", "dangles", "dscale", "dgamma"],
234
+ actual,
235
+ expected,
236
+ )):
232
237
  assert_close(f"compute_dqktheta_{name}", exp, act, 1e-3, err_atol=1e-3)
应用建议
likedislike
不准确？
233+ 
234+ 
235+@pytest.mark.parametrize(
236+ "batch,seqlen,nheads_qk,nheads,headdim_qk,dtype,chunk_size",
237+ TEST_CASES_DDT,
238+)
239+def test_compute_ddt_dtrap(batch, seqlen, nheads_qk, nheads, headdim_qk, dtype, chunk_size):
240+ if not hasattr(torch, "npu") or not torch.npu.is_available():
241+ pytest.skip("NPU is not available")
242+ device = torch.device("npu")
243+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk, dt, trap = make_inputs(
244+ batch, seqlen, nheads_qk, nheads, headdim_qk, dtype, device
245+ )
246+ 
247+ # First compute dqktheta to get dscale and dgamma
248+ dq, dk, dq_bias_out, dk_bias_out, dangles, dscale, dgamma = compute_dqktheta(
249+ q, k, scale, gamma, q_bias, k_bias, angles, dq_in, dk_in, dqk,
250+ chunk_size=chunk_size,
251+ )
252+ 
253+ actual_ddt, actual_dtrap, *_ = compute_ddt_dtrap_dinput_states(
254+ dscale, dgamma, dt, trap,
255+ )
256+ 
257+ expected_ddt, expected_dtrap = ref_compute_ddt_dtrap_dinput_states(
258+ dscale, dgamma, dt, trap,
259+ )
260+ 
261+ assert_close("compute_ddt_dtrap_ddt", expected_ddt, actual_ddt, 1e-2, err_atol=1e-2)
262+ assert_close("compute_ddt_dtrap_dtrap", expected_dtrap, actual_dtrap, 1e-2, err_atol=1e-2)