/label add good-first-issue

/label add triage-review

认领这个任务

认领这个任务

认领这个任务

认领这个任务

@qq_51696799

欢迎认领任务，请参考https://gitcode.com/Ascend/MindSpeed-Ops/issues/3 社区任务池明确该任务的：

• 完成的截止日期
• 开发进展反馈
• 微信答疑群
• 任务交付注意事项

等信息。如果您同时认领了多项任务，但无法都能进行投入，可以在部分任务中回复退出.

麻烦您加入到对应微信群，群备注名修改为"社区任务+您的gitcode账号", 后续有相关消息和问题都可以在微信群咨询答疑。 等您加入到微信群后，我这边会在社区任务池里面登记任务责任人。

已提交 PR !83：https://gitcode.com/Ascend/MindSpeed-Ops/merge_requests/83

已提交 PR !83：https://gitcode.com/Ascend/MindSpeed-Ops/merge_requests/83 @LinShua @Ronald1995

认领这个任务

@bitszh3271

你好，感谢您对昇腾平台的支持，本任务已有责任人认领，请关注仓上其他开放任务。

认领这个任务

@gcw_MTeIQk9o

你好，感谢您对昇腾平台的支持，本任务已有责任人认领，请关注仓上其他开放任务。

认领该任务

认领这个任务

mamba3_siso_fwd_kernel：Triton 算子 NPU 迁移实践与交付

交付 MR：Ascend/MindSpeed-Ops !194 — head b6098f764358c9734f5acae34e74b14c052d1c76，base master；当前 open，等待评审与合入。该 MR 携带本轮最新复现的源码与证据；此前的 !83 为较早 head，不含本轮精度修复。

交付内容（10 个文件）

• 生产源码：mindspeed_ops/api/triton/mamba3_siso_fwd_kernel.py、mindspeed_ops/arch32/triton/mamba3/{mamba3_siso_fwd_impl.py, mamba3_siso_fwd_kernels.py, mamba3_utils.py}
• 仓内 UT：tests/unit_tests/triton/{test_mamba3_siso_fwd_kernel.py, mamba3_siso_fwd_kernel_reference.py}
• ATK：tests/atk_tests/triton/mamba3_siso_fwd_kernel/{mamba3_siso_fwd_kernel.yaml, generate_mamba3_siso_fwd_kernel.py, triton_mamba3_siso_fwd_kernel.py}
• 文档：docs/triton/mamba3_siso_fwd_kernel.md

1. 算子介绍

mamba3_siso_fwd_kernel 实现 Mamba3 SISO dense、non-varlen forward 的主输出路径：对 Q/K 加 bias 并做 rotary 变换，由 adt 构造 chunk 内衰减，完成 QK-V 因果累积与跨 chunk 状态递推，并叠加可选 D skip 与 Z SiLU gating。

接口（详见 MR 内文档）：

mamba3_siso_fwd_kernel(q, k, v, adt, dt, trap, q_bias, k_bias, angles,
                       d=None, z=None, chunk_size=64) -> out  # out: [B, S, H, V]

输入：q/k: [B,S,Hq,K]、v: [B,S,H,V]、adt/dt/trap: [B,H,S]、q_bias/k_bias: [H,K]、angles: [B,S,H,A]（A <= K/2 且偶数）、d: [H]、z: [B,S,H,V]；Q/K/V/Z 支持 fp16 与 bf16，数值参数为 fp32；chunk_size ∈ {32, 64}；要求 H % Hq == 0。

生产 dispatch（只依赖运行时 shape/dtype/属性，无 fallback）：

arch35 未提供实现，调用时明确抛 NotImplementedError，不会切到 Torch/CPU/供应商整算子或其它后端。

2. 迁移要点（对比原 Triton 实现）

原 Triton 实现把 forward 拆为 scale / preprocess / q_rotary / k_rotary_scale / 主 recurrence 等阶段。本 PR 保留相同数学语义，并做以下 Ascend 适配与优化：

1. 将 CUDA/PTX 近似指令替换为 Triton JIT 内可编译的 sin/cos 多项式。
2. 对 single-chunk 常见路径新增全融合 kernel，减少中间张量与 launch。
3. 通用 dense 路径启用 INLINE_PREPROCESS，避免原 split 路径的多次预处理 launch。
4. 将 QK 点积改为显式逐行 reduction，降低矩阵乘 ones 带来的片上资源压力。
5. descriptor allocator 使用普通加速器内存分配，不依赖 NVIDIA TMA 的 alignment 约束；输入契约使用显式异常，不依赖可被优化模式移除的 assert。

原 split Triton 实现仅作为迁移来源，不复制进提交目录；生产 API 不导入测试代码。

3. 测试方法

环境：CANN 9.0.0 + Triton-Ascend；设备：本机物理 NPU 1（npu-smi 报告 Ascend910）。

# 仓内 UT
pip install -e .
pytest -q tests/unit_tests/triton/test_mamba3_siso_fwd_kernel.py

# ATK（26.9.24）：生成用例后分别跑 accuracy / performance_device / memory_device

4. 结果（绑定 head b6098f76...）

• UT：15 passed, 1 warning。覆盖 fp16/bf16 × 多组 shape/chunk_size、非连续输入、非法 chunk_size（非 int / bool / 非 32|64）、未请求 final states 不物化、allocator 对齐。精度基线是 UT-only 的 Torch 参考实现分别在 CPU 与 NPU 上执行（同一数学参考的两个设备实例，不是两个独立参考）。
• ATK accuracy：20/20 通过（Total Task: 20, success 20, failed 0，acc_pass_result: Pass）；判据为仓库 ATK 标准 single_bm（type: high_precision，fp16 误差阈值 2^-11、bf16 2^-8）。
• ATK performance_device：20/20 通过（device 性能通过率 100.0%，平均device性能比 59.3918，device性能是否达标: Pass）。该比值为 ATK 报告列 npu_0_Device性能 / triton_0_Device性能，即基线（ATK npu backend 上的 Torch 参考实现）设备耗时 ÷ 候选（triton backend）设备耗时，单位 x，为 20 例算术平均。
• ATK memory_device：20/20 通过（is_device_memory_pass: Pass）。
• 性能口径说明：ATK 性能基线是 NPU 上的 Torch 参考实现。任务书要求的“对比原 triton 算子提升 20%”若指上游 Triton 源码在 NPU 上的实现，本算子的本地 NPU-相对 active-window 实测为 geomean 12.40x、最小 6.04x、20/20 ≥ 1.2x；若指 CUDA 上的原 Triton 计时，则需 GPU 侧基准，请确认期望口径。

5. 已知边界

1. 仅交付 arch32；arch35 明确报不支持。
2. 公开 API 不覆盖 varlen、初始状态、state-return 与 backward 路径。
3. ATK 精度门全部通过。当输出量级接近 0（约 1e-9）时结果由 fp32 舍入主导，属该数值区间的固有精度地板；ATK 未对该区间单独设门，如需覆盖请确认判据。
4. 生成的 ATK 结果、日志、缓存与 profiler 目录不进入 MR diff。

mamba3_siso_fwd_kernel：Triton 算子 NPU 迁移实践与交付（issue 提交）

交付 MR：Ascend/MindSpeed-Ops !196 — head d261a24c0631bb85bc426da4905175c226e122eb，base master；当前 open，等待评审与合入。

交付内容：生产源码（mindspeed_ops/api/triton/、mindspeed_ops/arch32/triton/）+ 仓内 UT（tests/unit_tests/triton/）+ ATK 三件套（tests/atk_tests/triton/）+ 文档（docs/triton/mamba3_siso_fwd_kernel.md）。

验收证据（四项）

已在远端 head d261a24c0631bb85bc426da4905175c226e122eb（本 MR 当前 head，与复验工作树逐字节一致）上逐项复验，证据如下。

（1）ATK 通过截图

ATK 26.9.24，物理 NPU 1（Ascend910_9382），20 例：

ATK 三项（accuracy / performance_device / memory_device）在本 head 上全部通过；ATK 不在合并流水线门禁内（流水线只跑仓库原生 UT），此处作为附加验收证据。

（2）UT 通过截图

仓库原生 tests/unit_tests/triton/test_mamba3_siso_fwd_kernel.py：15 passed, 1 warning in 12.92s（物理 NPU 1，torch_npu 先于 triton 导入；即仓库 CI stage3 的合并门禁）。

（3）NPU 与 GPU 性能对比达标数据

口径：同一评测记录、同一窗口背靠背；任务书“相对原 Triton 算子提升 20%”的判据分母是 NPU 侧 triton_ascend_reference。

• 窗口：候选 eval_20260929_024106，参考刷新 eval_20260929_024106 (same record, back-to-back)（warmup 3 / repeat 5）。
• geomean(原 triton(NPU) / 候选) = 12.4032，min = 6.0442，≥1.2 覆盖 20/20。
• 与 GPU(H100 triton_cuda_reference) 的绝对对照：geomean(候选 / H100) = 2.3461（候选比 H100 慢约 2.3×，为 NPU/H100 硬件代差，不是任务书判据）。

（4）精度对齐

以 fp64 CPU golden（torch） 为共同高精度基准（独立于 NPU 与 GPU 两条执行路径）：

• NPU 候选（本 MR 算子）：19/20 通过；唯一未过为 case 16 的单个 bf16 元素（normal_mismatch_count=0，该 case 的 max_diff = 1.9073486e-6；20 例中 NPU 侧最大 max_diff 为 3.0517578e-5，出现在通过用例 case 12），判定为 framework fp32 累积地板（见 steward 判定文档，非候选缺陷、未放宽容差）。
• GPU(H100) triton 参考：max_diff 最大 1.2207031e-4（case 9）、mere=0、normal_mismatch_count=0，20/20 仅被 special-domain（相消）判据标记（special_domain_fail），无正常值失配。
• 结论：NPU triton 与 GPU triton 相对同一 fp64 golden 的绝对误差同量级（NPU ≤3.05e-5、GPU ≤1.22e-4）；两侧都只被 special-domain 判据标记，normal_mismatch_count=0。
• 逐 case 相对该 golden 的 max_diff/mere/判定（下表）：NPU 候选与 GPU triton 对同一高精度基准的误差同量级。
• 另：ATK accuracy（NPU 候选 vs 独立 NPU torch 小算子参考）20/20 Pass。

证据：截图存放在提交者 fork 的独立 evidence/pr196-mamba3-siso-fwd-atk 分支（commit e690ca58705140d963343d772bf914e0fb260e8b），不属于本 PR 的源码分支与交付文件 diff；ATK/UT 生成物与日志未提交到算子目录。

实践文档（PR 内 docs/triton/mamba3_siso_fwd_kernel.md 全文）

mamba3_siso_fwd_kernel 算子

概述

mamba3_siso_fwd_kernel 实现 Mamba3 SISO dense、non-varlen forward 的主输出路径。
算子对 Q/K 应用 bias 和 rotary 变换，根据 adt 构造 chunk 内衰减，完成
QK-V 因果累积与跨 chunk 状态递推，并可叠加 D skip 和 Z SiLU gating。

当前实现面向 arch32 NPU，使用 Triton-Ascend kernel。arch35 暂未提供
实现，调用时会明确抛出 NotImplementedError，不会切换到 Torch、CPU、
供应商整算子或其它后端。

模型接口

def mamba3_siso_fwd_kernel(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    adt: torch.Tensor,
    dt: torch.Tensor,
    trap: torch.Tensor,
    q_bias: torch.Tensor,
    k_bias: torch.Tensor,
    angles: torch.Tensor,
    d: torch.Tensor | None = None,
    z: torch.Tensor | None = None,
    chunk_size: int = 64,
) -> torch.Tensor:
    """Return the dense Mamba3 SISO forward output."""

输入说明

记 B 为 batch size，S 为 sequence length，Hq 为 Q/K head 数，H 为
输出 head 数，K 和 V 为两个 head dimension。H 必须能被 Hq 整除。

公开 API 通过仓内 input_guard 整理连续布局。生产 wrapper 显式检查 shape
代数、head 整除关系、rotary dimension 和 chunk_size；非法输入直接报错。

输出说明

实现原理

计算分为以下逻辑阶段：

1. 对 Q/K 加 bias，使用 angles 完成 rotary 变换。
2. 由 dt 和 sigmoid(trap) 计算 gamma、相邻 token scale 和 QK skip 项。
3. 在每个 chunk 内构造严格因果 QK score，并乘以 adt 推导的衰减。
4. 用上一 chunk 的状态贡献、chunk 内 QK-V 贡献和可选 D skip 生成输出。
5. 若传入 z，在同一生产路径中执行 SiLU gating；随后更新跨 chunk 状态。

生产调度只依赖运行时 shape、dtype、chunk_size 和是否启用可选输入：

当 K=128 且请求 chunk_size=64 时，wrapper 使用内部 32-token kernel tile
控制片上资源；该 tile 是实现细节，不改变公开输出 shape 或序列语义。所有
tail load/store 都按实际 S、K、V mask。

迁移要点

原始迁移实现将 forward 拆为以下 Triton 阶段：

• mamba3_siso_fwd_scale_kernel 生成 scale/gamma；
• mamba3_siso_fwd_preprocess_kernel 生成 Q/K/QK 中间量；
• mamba3_siso_fwd_q_rotary_kernel 和
  mamba3_siso_fwd_k_rotary_scale_kernel 分别处理 Q/K rotary；
• mamba3_siso_fwd_kernel 完成主 recurrence。

本 PR 的 candidate 保留相同数学语义，并做以下 Ascend 适配与优化：

1. 将 CUDA/PTX 近似指令替换为 Triton JIT 内可编译的 sin/cos 多项式。
2. 对 single-chunk 常见路径新增全融合 kernel，减少中间张量和 launch。
3. 对通用 dense 路径启用 INLINE_PREPROCESS，避免原 split 路径的多次预处理 launch。
4. 将 QK 点积改为显式逐行 reduction，降低矩阵乘 ones 带来的片上资源压力。
5. descriptor allocator 使用普通加速器内存分配，不依赖 NVIDIA TMA 的 alignment
   约束；输入契约使用显式异常，不依赖可被优化模式移除的 assert。

原 split Triton 实现作为迁移来源，不复制到提交目录；生产 API 不导入测试代码。

使用示例

import torch

from mindspeed_ops.api.triton.mamba3_siso_fwd_kernel import mamba3_siso_fwd_kernel

B, S, Hq, H, K, V = 1, 128, 2, 4, 64, 96
q = torch.randn((B, S, Hq, K), device="npu", dtype=torch.float16)
k = torch.randn_like(q)
v = torch.randn((B, S, H, V), device="npu", dtype=torch.float16)
adt = -torch.rand((B, H, S), device="npu", dtype=torch.float32) * 0.08
dt = torch.randn((B, H, S), device="npu", dtype=torch.float32) * 0.05
trap = torch.randn((B, H, S), device="npu", dtype=torch.float32)
q_bias = torch.randn((H, K), device="npu", dtype=torch.float16) * 0.05
k_bias = torch.randn_like(q_bias)
angles = torch.rand((B, S, H, K // 2), device="npu", dtype=torch.float32) * 0.25
d = torch.rand((H,), device="npu", dtype=torch.float32) * 0.2
z = torch.randn((B, S, H, V), device="npu", dtype=torch.float16)

out = mamba3_siso_fwd_kernel(
    q, k, v, adt, dt, trap, q_bias, k_bias, angles, d, z, chunk_size=64
)
print(out.shape)

适用场景

适用于 Mamba3 SISO dense、non-varlen、无初始状态、只返回主输出的 forward
路径，支持不同 batch、sequence、head、head dimension 以及序列尾 chunk。

注意事项

1. 当前生产交付仅支持 arch32；arch35 明确报不支持。
2. 当前公开 API 不覆盖 varlen、初始状态、state-return 或 backward 路径。
3. chunk_size 仅支持 32 和 64，H 必须能被 Hq 整除。
4. UT/ATK 中的 Torch 和原 Triton 实现只用于测试；生产路径不导入测试代码，
   不存在 Torch、CPU、供应商整算子、reference 或其它后端 fallback。

说明：上一条 192669818（作者 chongweiliu）是同一实践文档在 lcw 面的投递，其引用的 MR !194 已关闭、由 !196 取代；本条 192754987 是按账本账号的正式提交（MR !196、head d261a24c…），含四项验收证据与完整文档，请以本条为准。