变更摘要

本 PR 为 Mamba-3 MIMO 组合反向计算引入了第二遍 Triton-Ascend 内核 (mamba3_mimo_bwd_bwd_kernel) 的完整实现。该内核消费第一遍反向 (bwd_fwd) 缓存的 states 和 qk_dot 中间量，执行倒序逐 chunk 扫描，产出 12 个梯度（DQ、DK、DV、DMIMO_V、DANGLES、DFACTOR、DGAMMA、DDA、DSSDA、DDA_CS_REV、DDA_CS、DD）。实现采用三段式拆分架构（预处理/扫描/后处理），包含 baseline 参考内核和针对 arch32 的优化实现，支持 reduceO / 非 reduceO、可选 D-skip、GQA 等路径，并通过 host 端尾块填充 (pad_tail_seq / rebuild_dacs_tail) 处理非对齐序列长度。

主要改动

• 新增公开 API 接口：mindspeed_ops/api/triton/mamba3_mimo_bwd_bwd_kernel.py 定义 mamba3_mimo_bwd_bwd_kernel 函数，接收 19 个输入张量与 2 个属性参数，返回包含 12 个梯度的字典；内部根据架构（arch32/arch35）路由到对应实现。

• 新增三段式 Triton baseline 内核：mamba3_mimo_bwd_bwd_baseline_impl.py 实现 _mimo_bwd_pre_kernel（旋转编码 + PsiV 投影 + dPhiO）、_mimo_bwd_scan_kernel（倒序 chunk 扫描、状态反传、逐 rank 块对角 matmul）、_mimo_bwd_post_kernel（逆旋转 + 对角修正 + dangle），全程 fp32 累加，支持最小可测路径。

• 新增优化 Triton-Ascend 实现：mamba3_mimo_bwd_bwd_impl.py（+1479 行）包含针对 arch32 的优化版本，引入 R4 rank-fold 优化、N-blocked 反向扫描、运行时 rank-loop 选择等策略，将 chunk 局部矩阵运算从串行反向扫描中分离以提升性能。

• 新增尾块填充工具函数：mindspeed_ops/api/triton/utils.py 新增 pad_tail_seq（支持 zero/edge 两种填充模式）和 rebuild_dacs_tail（填充后重建尾部 da_cs_rev 与 segsum），为 host 端处理非 chunk 对齐序列长度提供基础支撑。

• 新增完整测试体系：包含 ATK 配置 (mamba3_mimo_bwd_bwd_kernel.yaml)、用例生成器 (generate_mamba3_mimo_bwd_bwd_kernel.py)、torch 参考实现 (reference_impl.py，含前向缓存复算 _fwd_cache 与反向参考 _bwd_ref)、ATK API 注册 (triton_mamba3_mimo_bwd_bwd_kernel.py) 以及单元测试 (test_mamba3_mimo_bwd_bwd_kernel.py)，覆盖基础参数网格、官方 TileLang 形状网格（quick/full）和尾块非整除路径。

代码审查

Now let me provide the closing summary.

---

审查总结

逐文件审查结果：

文件	审查结论
docs/triton/mamba3_mimo_bwd_bwd.md	无问题 — 纯文档，无危险命令/密钥/不安全配置
docs/triton/mamba3_mimo_bwd_bwd_optimization.md	无问题 — 纯文档
docs/triton/mamba3_mimo_bwd_bwd_practice.md	无问题 — 纯文档
mindspeed_ops/api/triton/mamba3_mimo_bwd_bwd_kernel.py	无问题 — 公开 API，架构检查与委托正确
mindspeed_ops/api/triton/utils.py	无问题 — pad_tail_seq/rebuild_dacs_tail 逻辑正确
mindspeed_ops/arch32/triton/mamba3/mamba3_mimo_bwd_bwd_baseline_impl.py	无问题 — 基线实现，作为对照保留
mindspeed_ops/arch32/triton/mamba3/mamba3_mimo_bwd_bwd_impl.py	2 个问题：P2 BN 搜索非整除 + P3 scan mask 死代码
tests/atk_tests/.../generate_mamba3_mimo_bwd_bwd_kernel.py	1 个问题：P3 生成器 z/mimo_z 配置不一致
tests/atk_tests/.../mamba3_mimo_bwd_bwd_kernel.yaml	无问题 — YAML 键名无拼写错误
tests/atk_tests/.../reference_impl.py	1 个问题：P3 _fwd_cache 中 causal 死代码
tests/atk_tests/.../triton_mamba3_mimo_bwd_bwd_kernel.py	无问题 — ATK 适配代码正确
tests/unit_tests/triton/test_mamba3_mimo_bwd_bwd_kernel.py	无问题 — 测试覆盖充分

问题按优先级统计：

• P0: 0
• P1: 0
• P2: 1（BN 块大小非整除的潜在正确性问题，confidence 0.65）
• P3: 3（死代码与测试配置不一致，均为可选的清理改进）

整体风险判断：低。 核心数学逻辑（旋转梯度、递推状态传播、rank-fold 矩阵收缩、尾块填充/裁剪）经过仔细验证均正确。P2 问题仅在非标准 N 值（非 2 的幂）且 N*P > 8192 时触发，当前所有测试 shape 均不涉及，但建议修复以避免未来风险。三个 P3 问题均为代码清洁度问题，不影响运行时正确性。

⚠️ 已识别出整体风险，但无法提取行内评论，请参考整体评估。

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.

---

PR Approval Progress

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

Module Approval Details

module	lgtm status	approve status
docs	❌ LinShua (1/2)(You can also ask: zhizaidicengshehua, 雷珍珍, 刘荣, chic666, guofanfeng23)	❌ (0/1)(You can also ask: 朱彦儒, 王晓歆, 华郁秀, 刘荣, 雷珍珍)
repo-Ascend/MindSpeed-Ops	❌ LinShua (1/2)(You can also ask: 华郁秀, guihaowen666, zjchenn, feng0w0, 刘荣)	❌ (0/1)(You can also ask: 周蓓蓉, 华郁秀, 刘哲续, bigdog1206, 刘荣)

💡 Tip:

• Committer can comment /approve or /lgtm
• Commenting /approve implies both code review (lgtm) and intent to merge (approve)

---

CLA Signature Pass

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

mamba3_mimo_bwd_bwd Triton-Ascend 实践

背景

mamba3_mimo_bwd_bwd 完成 Mamba-3 MIMO combined backward 的第二次扫描。第一次扫描沿正向
重算并保存 chunk 进入状态；第二次扫描从最后一个 chunk 开始维护 dstates，计算 q、k、v、Psi、
旋转角和离散化链的梯度。

上游 GPU TileLang 实现将大量 chunk-local 小矩阵运算放在倒序循环内。直接迁到 Ascend 可以得到
正确结果，但串行路径会被 tiny-dot 启动、block pointer 生成和 rank tile 重载主导。本实践文档说明
如何在不改变缓存契约的前提下拆分这些工作。稳定接口见 mamba3_mimo_bwd_bwd.md，优化收益见
mamba3_mimo_bwd_bwd_optimization.md。

输入缓存的语义

states

states[b,h,c] 是第 c 个 chunk 开始前的前向状态，而不是该 chunk 结束后的状态。倒序扫描需要
它计算当前 q 对历史状态的贡献，并通过 chunk 衰减将 dstates 传到前一个 chunk：

state_out = state_in * exp(da_cs_sum) + K_state.T @ PsiV

dstates_in = dstates_out * exp(da_cs_sum)
             + Q.T @ (dPhiO * exp(da_cs))

这里省略了 rank 展平和 chunk 内因果项。缓存位置一旦错一位，短序列可能仍能产生有限数值，但所有
跨 chunk 梯度都会失去对应的前向状态。

qk_dot

qk_dot[b,h,s,r_out,r_in] 保存同一 token 上各 rank 的 q-k 点积。前向对角项使用
gamma * qk_dot * PsiV，第二阶段据此恢复 DV、DGAMMA 以及 q/k 的对角修正。该缓存只包含同位置
的 R x R 块；chunk 内严格因果的 q-k 矩阵仍在本阶段由旋转后的 q/k 计算。

反向语义拆解

上游梯度映射

reduce-O 路径先把 [B,S,H,P] 的 dout 展到 rank：

dPhiO[s,r,p] = dout[s,p] * mimo_o[r,p]

非 reduce-O 路径直接接收 [B,S,R,H,P]。两条路径从此共享相同的 dPhiO[C,R,P] 表示。
value 投影中间量为：

PsiV[s,r,p] = v[s,p] * mimo_v[r,p]

pre kernel 同时生成这两个张量，后续 aux 与 scan 不再区分上游布局。

梯度分组

12 项输出可按数据依赖分为三组：

分组	梯度	依赖
chunk-local	DGAMMA、DSSDA、部分 DQ/DK/DV	当前 chunk 的 q/k、PsiV、dPhiO、qk_dot
state chain	DMIMO_V、DDA、DDA_CS、DDA_CS_REV、部分 DK/DV、DD	states 与倒序传入的 dstates
rotary post	最终 DQ/DK、DANGLES	旋转前 q/k、旋转域梯度与 angles

第一组可以在 chunk 维并行，第二组必须按 chunk 倒序，第三组在旋转域梯度完成后再次并行。
生产实现的 pre/aux/scan/post 切分正是沿这三类依赖建立的。

从基线到生产实现

stage0

mamba3_mimo_bwd_bwd_baseline_impl.py 保留优化前版本。每个 (batch, head) program 在一个倒序
chunk 循环中完成：

1. bias-add、旋转和离散化量构造；
2. q-k、q-state、PsiV-dPhiO 等局部矩阵运算；
3. dstates 更新和 Psi/D 梯度累加；
4. 逆旋转、角度梯度和 12 项输出写回。

该版本便于逐公式核对，但所有工作都受倒序依赖串行化，也让单个 program 的 UB 和编译规模随
N、P、R 一起增长。

生产 kernel 切分

kernel	任务粒度	作用
_mimo_bwd_pre_kernel	(B,H,nchunks)	q/k 旋转、PsiV、dPhiO、离散化 scratch
_mimo_bwd_aux_kernel	(B,H,nchunks)	通用 rank 的 chunk-local 矩阵运算
_mimo_bwd_aux_kernel_rf	(B,H,nchunks)	R4 的 rank-fold chunk-local 运算
_mimo_bwd_scan_kernel	(B,H)	状态常驻的倒序递推
_mimo_bwd_scan_kernel_bt	(B,H)	大 N/P 的 N-blocked 倒序递推
_mimo_bwd_post_kernel	(B,H,nchunks)	逆旋转、q/k 对角修正和 DANGLES

pre 写出 Qrot/Krot 供矩阵运算使用，同时保留 Qpre/Kpre 供角度梯度使用。aux 写出
DQcur/DKLOC/DVLOC/DQKG 等 scratch。scan 在读取这些局部结果后，只完成与 dstates 有关的
收缩、状态传递和跨 chunk 累加。post 最后把旋转域梯度映射回输入布局。

Ascend 侧实现要点

将可并行矩阵乘移出 scan

倒序依赖只约束 dstates，并不约束当前 chunk 内的 q-k 或 PsiV-dPhiO 矩阵。aux 使用
B*H*nchunks 网格后，这些运算可以在不同 AI Core 上同时执行，并由 Triton-Ascend 调度到 Cube。
scan 与 aux 之间通过 fp32 scratch 交换局部结果，以额外显存换取串行关键路径缩短。

这种拆分的边界必须保持清楚：凡是读取当前 dstates 的计算仍在 scan；只依赖 states[c] 和当前
chunk 输入的计算可以放到 aux。把前者误移到 aux 会破坏倒序递推，而把后者留在 scan 只会增加
串行时间。

R4 rank-fold

通用 aux 按 (r_out, r_in) 处理多个 C x C 小矩阵。R4 时，rank 与 chunk 折为 R*C，构造：

Qf, KSCf : [R*C, N]
PSIVf, DPHIOf : [R*C, P]

随后以少量较大矩阵完成 q-state、k-q、PsiV-dPhiO、DV 和 DK 的主要收缩。Triton-Ascend 的
二维 gather load 必须显式提供 mask；即使 shape 在逻辑上完整，也不能依赖无 mask 的偏移寻址。

R2 的矩阵规模过小，折叠的数据整理成本高于收益；R8 的 [R*C,*] tile 超出 UB 预算。因此生产
调度只对 R4 使用 rank-fold，其他 rank 走通用 aux。

dstates 的 UB 管理

常规 scan 让 [N,P] 的 dstates 常驻 UB，可以避免每个 chunk 访问 HBM。当 N*P > 8192 时，
前向 states[c]、dstates 与 q/k/PsiV tile 的总占用超过预算，切换为 N-blocked scan：

for n_block in N / BN:
    load dstates[BN,P] from HBM scratch
    reverse scan all chunks for this block
    store dstates[BN,P]

BN 取能整除 N 且不超过 64 的值。分块增加状态 scratch 的 HBM 往返，但使 UB 占用不再随完整 N
增长，并让大 shape 获得可编译、可运行的实现。

rank 循环的编译策略

通用 aux 和 post 含有多处 R x R 收缩。R6/R8 全部静态展开会生成很大的 IR；大 tile 下即使 R4
也可能出现同样问题。host 在 R >= 6 或 N*P > 8192 时使用运行时 rank 循环，其余 shape 保留
static_range。该策略把编译规模控制在可接受范围，同时避免所有 shape 都支付运行时循环开销。

旋转梯度后处理

q/k 的矩阵运算在旋转域完成。post 使用保存的旋转前分量计算：

dangle = dq_rot * d(q_rot)/d(angle) + dk_rot * d(k_rot)/d(angle)

再应用旋转矩阵的转置得到输入 q/k 梯度，并加上 qk_dot 对角项的修正。把这一阶段独立出来后，
角度梯度的 rank 求和和序列写回不占用倒序 scan 的执行时间。

非整 chunk

公开入口在 host 侧把序列补到完整 chunk。q/k/v/dout/angles 和 qk_dot 尾部补零，da_cs 用末值
延展，并重建 da_cs_rev/segsum。执行后按原 S 裁剪 DQ、DK、DV、DANGLES 和离散化梯度；
DSSDA 最后一个 chunk 的无效行列显式清零。states 的 chunk 数在补齐前后相同，可直接复用。

数值与验证方法

生产实现的中间 scratch 与归约使用 fp32。output_dtype 只控制 DQ、DK、DV 的最终写回，其余梯度
保持 fp32。由于 aux 拆分和 rank-fold 改变了部分矩阵的结合顺序，验证不要求逐位一致，而是与独立
fp64 golden 比较。

验证分为三层：

1. PyTorch 小算子重算前向缓存，并以显式公式执行倒序扫描；
2. stage0 与 reference 对齐，确定迁移前 Triton 语义；
3. 生产实现覆盖 pairwise、上游 quick/full shape 和非整 chunk，用 fp32/fp64 dual-benchmark
   逐项检查 12 个梯度。

DDA 在个别 fp16 shape 上会出现 fp32 reference RMSE 接近零的情况，DANGLES 也可能含有接近
零的角度梯度。候选即使整体更接近 fp64 golden，也可能因误差比值的分母过小而失败。这两项使用
候选对 fp64 golden 的全局最大相对误差，直接约束绝对误差和整体相对误差。

ATK 使用相同 reference 与候选入口验证标准 shape。生产代码不导入测试 reference，不支持的 arch35
直接报错，不使用其它整算子作为运行时替代。

扩展实现时的检查项

• 修改第一阶段缓存格式时，同时核对 states 的进入态语义和 qk_dot 的 rank 顺序。
• 新增梯度分支时，先判断它属于 chunk-local、state chain 还是 rotary post，再选择 kernel 阶段。
• 修改 aux scratch 时确认所有元素均在 scan 读取前写入，避免 torch.empty 引入未初始化值。
• 调整 rank-fold 时覆盖 R1/R2/R4/R8，并检查二维 gather 的 mask 与 UB 总量。
• 调整 N-block 时验证 BN 整除 N，并对比常驻 scan 的浮点结果。
• 修改尾块逻辑时逐项检查 12 个输出的裁剪规则，特别是 DQ/DK 的 S*R 重排和 DSSDA 清零。
• 性能评估使用仓内 stage0；stage0 因 UB 无法运行的 shape 单独报告，不纳入加速比统计。

/check-cla

mamba3_mimo_bwd_bwd 算子迁移说明

算子概述

mamba3_mimo_bwd_bwd_kernel 是 Mamba-3 MIMO combined backward 的第二阶段。算子接收第一阶段
mamba3_mimo_bwd_fwd_kernel 生成的 states 和 qk_dot，按 chunk 倒序传播状态梯度，输出
q、k、v、Psi、旋转角和离散化链上的 12 项梯度。

当前实现面向 Ascend arch32，公开入口位于
mindspeed_ops.api.triton.mamba3_mimo_bwd_bwd_kernel。arch35 暂未实现，调用时会明确抛出
NotImplementedError。

函数签名

def mamba3_mimo_bwd_bwd_kernel(
    dout,
    q,
    k,
    v,
    q_bias,
    k_bias,
    mimo_v,
    states,
    qk_dot,
    da_cs,
    da_cs_rev,
    dt,
    trap,
    segsum,
    angles,
    mimo_o=None,
    d=None,
    chunk_size=16,
    rotary_dim_divisor=4,
    output_dtype=None,
):
    ...

输入与输出

设 batch size 为 B，序列长度为 S，MIMO rank 为 R，KV head 数为 G，query/key
维度为 N，attention head 数为 H，value 维度为 P，chunk 大小为 C。

参数	形状	说明
dout	[B,S,H,P] 或 [B,S,R,H,P]	reduce-O 或逐 rank 的上游梯度
q, k	[B,S,R,G,N]	query、key
v	[B,S,H,P]	value
q_bias, k_bias	[H,R,N]	旋转位置编码前的偏置
mimo_v	[H,R,P]	value 投影参数 Psi
states	[B,H,ceil(S/C),N,P]	第一阶段保存的 chunk 起始状态
qk_dot	[B,H,S,R,R]	第一阶段保存的同位置 q-k 点积
da_cs, da_cs_rev, dt, trap	[B,H,S]	状态离散化中间量
segsum	[B,H,ceil(S/C),C,C]	chunk 内段和
angles	[B,S,H,N / rotary_dim_divisor]	旋转位置编码角度
mimo_o	[H,R,P] 或 None	output 投影参数 Phi；非空时启用 reduce-O
d	[H] 或 None	可选 D-skip 参数

返回值为包含以下键的字典：

键	形状	dtype 与含义
DQ, DK	[B,S*R,H,N]	q、k 梯度；默认跟随各自输入 dtype
DV	[B,S,H,P]	v 梯度；默认跟随 v dtype
DMIMO_V	[B,H,R,P]	Psi 梯度，fp32；调用方按 batch 维求和
DANGLES	[B,S,H,RD]	旋转角梯度，fp32
DFACTOR	[B,H,S]	trap_scale 链路梯度，fp32
DGAMMA	[B,H,S]	对角项 gamma 梯度，fp32
DDA	[B,H,S]	chunk 状态衰减梯度，fp32
DSSDA	[B,H,ceil(S/C),C,C]	chunk 内 segsum 梯度，fp32
DDA_CS_REV	[B,H,S]	反向累计衰减量梯度，fp32
DDA_CS	[B,H,S]	正向累计衰减量梯度，fp32
DD	[B,H] 或 None	D-skip 梯度，fp32

显式传入 output_dtype 时，仅覆盖 DQ、DK 和 DV 的输出 dtype；其余输出保持 fp32。

在 combined backward 中的位置

bwd_fwd
  ├── states：每个 chunk 的前向起始状态
  └── qk_dot：每个 token 的 R x R 对角点积
             │
             ▼
          bwd_bwd
  pre ──► aux ──► reverse scan ──► post
                                └── 12 项梯度

倒序扫描维护 dstates[N,P]。对当前 chunk，先将上游梯度映射到每个 rank：

dPhiO[s,r,p] = dout[s,p] * Phi[r,p]       # reduce-O
dPhiO[s,r,p] = dout[s,r,p]                # 非 reduce-O
PsiV[s,r,p]  = v[s,p] * Psi[r,p]

chunk 内局部 q-k 交互产生 q、k、v、Psi 和 segsum 的梯度；跨 chunk 路径使用 states、
da_cs、da_cs_rev 更新当前 dstates，再传递到前一个 chunk。qk_dot 用于恢复前向中
同位置 rank 交互的对角修正，无需在第二阶段重复计算。

实现说明

实现拆为四类 kernel，使串行部分只保留真正的状态依赖：

阶段	并行粒度	主要工作
pre	(B,H,nchunks)	bias、rotate-half、PsiV、dPhiO，并预计算 gamma、指数衰减和因果 mask
aux	(B,H,nchunks)	chunk-local 的 q-k、q-state、v/Psi、DSSDA 等矩阵收缩
reverse scan	(B,H)	倒序更新 dstates，汇总跨 chunk 的 DV、DK、Psi、D 与离散化梯度
post	(B,H,nchunks)	逆旋转 q/k 梯度、对角修正并汇总 DANGLES

pre 阶段保留旋转前的 Qpre/Kpre 和旋转后的 Qrot/Krot。前者供 post 计算角度梯度，后者供
aux 与 scan 使用。所有中间点积与归约使用 fp32，避免低精度 scratch 对长递推链造成额外误差。

aux 阶段先于 reverse scan 执行，并将 DVLOC、DKLOC、DQcur 等局部结果写入 scratch。
reverse scan 读取这些结果后只需完成与 dstates 有关的点积和状态传递，避免把可并行的工作留在
串行关键路径上。

当 N * P > 8192 时，完整 dstates tile 与其它中间量会超过 UB 预算。实现改用 N-blocked scan，
将 dstates 放入 HBM scratch 并按 BN <= 64 流式更新；其余 shape 使用状态常驻的常规 scan。

支持范围

• 硬件：Ascend arch32。
• dtype：float16、bfloat16、float32。
• 布局：dense、非 varlen，H % G == 0。
• 支持 reduce-O 与逐 rank dout，由 mimo_o 是否为空决定。
• 支持可选 D-skip、通用 rank 和非整 chunk 的序列长度；已覆盖 R = 1, 2, 4, 8。
• N 必须为偶数，且旋转维度不得超过 N / 2。
• 当前公开 API 不包含 cu_seqlens、Z 门控和 fuse_pregate_headwise_rms_norm；门控对
  dout 的影响应由第一阶段处理完成。

当 S % chunk_size != 0 时，host wrapper 对序列输入补零、对 da_cs 使用末值延展，并重建
尾块的 da_cs_rev 与 segsum。qk_dot 的尾部补零；执行后裁剪所有序列梯度，并将
DSSDA 最后一个 chunk 的无效行列清零。

调用示例

from mindspeed_ops.api.triton.mamba3_mimo_bwd_bwd_kernel import (
    mamba3_mimo_bwd_bwd_kernel,
)

grads = mamba3_mimo_bwd_bwd_kernel(
    dout,
    q,
    k,
    v,
    q_bias,
    k_bias,
    mimo_v,
    states,
    qk_dot,
    da_cs,
    da_cs_rev,
    dt,
    trap,
    segsum,
    angles,
    mimo_o=mimo_o,
    d=d,
    chunk_size=16,
    output_dtype=torch.float32,
)

dq = grads["DQ"]
dk = grads["DK"]
dv = grads["DV"]

代码与测试

文件	内容
mindspeed_ops/api/triton/mamba3_mimo_bwd_bwd_kernel.py	公开 API 与架构检查
mindspeed_ops/arch32/triton/mamba3/mamba3_mimo_bwd_bwd_baseline_impl.py	优化前 Triton 基线
mindspeed_ops/arch32/triton/mamba3/mamba3_mimo_bwd_bwd_impl.py	优化后的生产实现
tests/atk_tests/triton/mamba3_mimo_bwd_bwd_kernel/reference_impl.py	PyTorch 小算子参考实现
tests/unit_tests/triton/test_mamba3_mimo_bwd_bwd_kernel.py	双标杆、官方 shape 网格和尾块测试
tests/atk_tests/triton/mamba3_mimo_bwd_bwd_kernel/	ATK 用例配置与适配代码

PyTorch 参考先重算 states 与 qk_dot，再用独立的小算子公式执行倒序扫描。测试分别用 fp32
reference 和 fp64 golden 计算 12 项梯度，并要求 Triton 相对 golden 的误差不劣于同精度 reference。
DDA 在部分 fp16 shape 上的 fp32 reference 误差接近零，DANGLES 也可能包含接近零的角度梯度；
直接计算逐元素误差比值会放大误差底噪。这两项改用
max(abs(candidate-golden)) / max(abs(golden)) 的全局相对误差。

用例组	覆盖内容
pairwise 常规用例	B={1,2}、S={32,48}、H={2,4}、fp16/bf16/fp32、reduce-O、D、GQA
official quick	5 组代表 shape，覆盖 R1/R2/R4、N=16..256、P=64..128
official full	B=4, S=2048, H=16, G=1 的 11 组完整网格，包含 R8；由环境变量开启
tail	S={33,40,50}、C=16，校验输入右填、缓存补齐、12 项输出裁剪

pytest -q tests/unit_tests/triton/test_mamba3_mimo_bwd_bwd_kernel.py

完整 shape 网格通过 MIMO23_RUN_OFFICIAL_FULL_GRID=1 开启。ATK 精度任务覆盖 16 个用例，
当前记录为 16/16 通过。

mamba3_mimo_bwd_bwd 优化记录

基线与目标

优化基线保存在
mindspeed_ops/arch32/triton/mamba3/mamba3_mimo_bwd_bwd_baseline_impl.py。stage0 按上游
TileLang 的计算顺序执行倒序扫描：每个 (B,H) program 在 chunk 循环内完成离散化、局部矩阵运算、
状态梯度更新和最终梯度写回。

生产实现位于 mamba3_mimo_bwd_bwd_impl.py。优化保持输入、12 项输出、fp32 累加位置和
states/qk_dot 缓存契约不变，目标是把不依赖 dstates 的工作移出串行关键路径，并解决大 tile、
大 rank 的编译和 UB 限制。

瓶颈分析

stage0 的主要问题不是矩阵乘法计算量本身，而是串行 scan 内大量小矩阵的调度开销。每个 chunk
需要反复建立 block pointer、加载 rank tile 并启动 C x C 或更小的点积。profile 中 scan 的
Cube 利用率低，标量地址生成占比较高；随着 S/C 增长，这部分固定开销沿串行循环累积。

初步将离散化量移出 scan 后，热点转移到 chunk-local 矩阵收缩。进一步 profile 显示 aux kernel
约占生产路径 wall time 的 65%，仍然是 scalar-bound：R4 下每个 chunk 需要约
R + 7 * R * R 组 tiny-dot，矩阵计算规模不足以摊薄启动和寻址成本。

此外还有两类 shape 边界：

• N * P > 8192 时，dstates、前向状态及 q/k tile 同时驻留会超过 UB 预算；
• R6/R8 或大 tile 下静态展开 R x R 循环会放大编译期 IR，导致编译时间和 cbuf 需求失控。

优化项与对应收益

1. 离散化预计算

将只依赖当前 chunk 输入的 gamma、trap_scale、exp(da_cs)、exp(da_cs_rev)、chunk 衰减和
因果 mask 移到 pre kernel。pre 按 (B,H,nchunks) 并行，reverse scan 只读取预计算结果。

这一步把每个 chunk 内重复执行的指数与 mask 构造从串行循环中删除，并为后续 aux 拆分提供统一的
fp32 scratch。相较 stage0，scan 内的指数计算由多次降为零。

2. chunk-local 计算迁移到 aux

将以下不依赖 dstates 的计算移到并行 aux kernel：

• DGAMMA 与 q-k 对角修正 DQKG；
• DDA_CS 及局部 DQcur；
• DVLOC、DKLOC 与 DSSDA；
• chunk 内 q-k、PsiV 与 dPhiO 的矩阵收缩。

aux 使用 (B,H,nchunks) 网格，Triton-Ascend 可将这些矩阵运算调度到 Cube。reverse scan 最终只
保留 dstates 相关的 dp_state、dki_raw 点积和跨 chunk 递推。一次代表性 profile 中，串行
scan 设备时间从 5274 us 降至 365 us；整体 stage0 加速来自这项结构拆分，而非单个 load 微调。

3. aux rank-fold

R4 仍受大量微矩阵启动限制。生产实现将 rank 与 chunk 维按 r * C + c 折叠，构造
Qf/KSCf/PSIVf/DPHIOf[R*C,*]，用 6 个较大的矩阵收缩完成：

dqinter, dk_all, lkq_all, dq_comb, dvl, dkl

DGAMMA、DQKG 和 qk 对角修正保留向量实现。R4 下这一改动用少量大矩阵替代约 116 个
tiny-dot，aux 阶段额外获得约 1.6x–2.1x 加速。2D gather load 显式使用 mask，保证尾部地址
不会参与计算。

rank-fold 只对 R == 4 启用：R2 的折叠收益不足以覆盖额外数据整理，R8 的 folded tile 超出 UB
预算，R1 没有可折叠的 rank 组合。这些 shape 保留通用 aux 路径。

4. N-blocked reverse scan

当 N * P > 8192 时，将 dstates[N,P] 放入 HBM scratch，按不超过 64 的 BN 分块读取和更新。
该路径把 UB 峰值从 O(N*P) 限制为 O(BN*P)，覆盖 N256_P64、N128_P128 等大 tile。
阈值以内仍使用状态常驻版本，避免分块带来的 HBM 往返。

5. rank 循环 de-unroll

R6/R8 和大 tile 对 static_range 展开非常敏感。host 在 R >= 6 或 N * P > 8192 时改用运行时
rank 循环，将多个 R x R 静态展开降为单层循环。浮点运算顺序保持不变，小 rank 则继续使用静态
展开以避免约 5% 的串行开销。

6. 独立 post 与输出初始化

逆旋转及 DANGLES 汇总移到并行 post kernel，避免增加 reverse scan 的串行长度。生产实现确认
DANGLES、DFACTOR、DGAMMA、DDA、DSSDA、DDA_CS_REV、DDA_CS 均由 kernel
完整覆盖，因此使用 torch.empty 代替 torch.zeros，减少 7 次 memset launch；该项在代表 shape
上约带来 0.9% 收益。

性能结果

常规 shape

以下结果来自 Ascend 910B 单卡、fp16、同步 host wall-clock 中位数。生产实现与 stage0 使用相同
输入和输出，PyTorch eager 仅作为端到端参考：

Shape (B,S,H,G)	PyTorch eager	stage0	生产实现	相对 eager	相对 stage0
1,64,4,2	104.389 ms	4.620 ms	1.784 ms	58.5x	2.59x
1,512,8,4	1611.442 ms	30.447 ms	1.984 ms	812.4x	15.35x
8,512,16,8	30161.441 ms	213.657 ms	13.963 ms	2160.0x	15.30x

短序列中 kernel launch 和 scratch 往返占比更高，因此加速比低于中长序列；随着 chunk 数增加，
从串行 scan 移出的工作增多，pre/aux 并行化收益更明显。

上游 11-shape 网格

以下结果使用 B=4, S=2048, H=16, G=1、bf16，在同卡交错执行候选与 stage0：

Shape (N_P_R_C)	相对 stage0
N16_P64_R4_C8	13.52x
N32_P64_R4_C16	29.45x
N64_P64_R4_C16	31.55x
N128_P64_R4_C16	34.81x
N64_P128_R4_C16	29.62x
N128_P32_R4_C16	36.05x
N128_P64_R8_C8	8.16x

7 个 stage0 可运行 shape 的几何平均加速比约为 23.4x。其中 R4 rank-fold 相比折叠前的 aux
实现继续提升 1.6x–2.1x。N256_P64_R4_C8、N128_P128_R4_C8、
N128_P64_R2_C32、N128_P64_R1_C64 的朴素 stage0 本身因 UB 溢出无法运行，候选实现可运行，
但这 4 项不进入加速比统计。

精度与边界验证

• UT 使用 PyTorch 小算子分别生成 fp32 reference 和 fp64 golden，逐项检查 12 个梯度。
• DDA 的 fp32 reference 在个别 fp16 shape 上接近 fp64 golden 的误差底噪，DANGLES 也包含
  接近零的角度梯度；两项使用候选对 golden 的全局最大相对误差。
• 常规 pairwise 用例覆盖 fp16/bf16/fp32、reduce-O/逐 rank、D-skip 和 GQA。
• official quick 与 full 网格覆盖 R={1,2,4,8}、N=16..256、P=32..128、
  C={8,16,32,64}。
• rank-fold 引入的浮点重结合误差最大相对值记录不超过 1.6e-7，完整验证集的最大相对误差不超过
  7e-3。
• tail 用例验证输入右填、qk_dot 补零、所有序列输出裁剪以及 DSSDA 无效区域清零。
• ATK 精度任务为 16/16 通过。

性能比较只使用相同语义的 Triton stage0；无法运行的基线 shape 明确标记并排除，不以 PyTorch
reference 或其它整算子替代基线。

compile

阶段	任务名	状态	详情
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	Only_doc_commit	✅	>>>
	pre-commit	❌	>>>
开源片段检查	SCA	❌	>>>
开发者测试	UT_MindSpeed	🕚	>>>
流水线	PR-pipeline_MindSpeed-Ops	❌	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, liuzhexu]评论有效

• compile : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

compile

阶段	任务名	状态	详情
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	Only_doc_commit	✅	>>>
	pre-commit	❌	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_MindSpeed	🕚	>>>
流水线	PR-pipeline_MindSpeed-Ops	❌	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, liuzhexu]评论有效

• compile : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

compile

阶段	任务名	状态	详情
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	Only_doc_commit	✅	>>>
	pre-commit	✅	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_MindSpeed	❌	>>>
流水线	PR-pipeline_MindSpeed-Ops	❌	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, liuzhexu]评论有效

• compile : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

compile

阶段	任务名	状态	详情
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	Only_doc_commit	✅	>>>
	pre-commit	✅	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_MindSpeed	✅	>>>
流水线	PR-pipeline_MindSpeed-Ops	✅	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, liuzhexu]评论有效

• compile : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

330ea196f3114679bccd47ff2d861ee0.xlsx
atk精度报告

@LinShua 已经更新

您好，当前实现的NPU性能与GPU的性能差异较大，暂无法合入，麻烦进一步优化相关性能