MindSpeed Ops

简介

提供昇腾优化的训练业务自定义算子实现

详细文档请参见docs/zh，包括简介、软件安装、快速入门、算子清单、版本说明与安全声明。

版本说明

软件	版本
MindSpeed Ops分支	master
CANN版本	9.1.0
PyTorch	>=2.7.1
triton-ascend	3.2.2
Python版本	>=Python3.10.x

安装

安装依赖的软件

在安装MindSpeed Ops之前，请参考版本配套表，安装配套的昇腾软件栈，软件列表如下：

依赖软件	软件安装指南
昇腾NPU驱动	《驱动固件安装指南》
昇腾NPU固件
Toolkit（开发套件）	《CANN 软件安装指南》
Kernel（算子包）
NNAL（Ascend Transformer Boost加速库）
PyTorch	《Ascend Extension for PyTorch 配置与安装》
torch_npu插件
apex

获取MindSpeed Ops源码并安装

git clone https://gitcode.com/Ascend/MindSpeed-Ops.git

cd MindSpeed-Ops
# 如果原先有mindspeed_ops安装，请先卸载再重新安装
pip install -e . --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple --no-build-isolation --no-deps

可选：编译并运行 CATLASS 算子

MindSpeed Ops 默认不编译 CATLASS 算子。CATLASS 算子当前仅支持 ascend910b 与 ascend910_93 系列，Ascend 950 系列暂不支持。 在已安装 CANN、PyTorch 和 torch_npu 的环境中，可以启用 CATLASS 编译：

source /usr/local/Ascend/ascend-toolkit/latest/set_env.sh
export MINDSPEED_BUILD_CATLASS=1

pip install -e . --no-build-isolation --no-deps

setup.py 默认通过 npu-smi 自动识别芯片。容器内无法使用 npu-smi 时，可以显式设置：

# Atlas A3 示例；ascend910b* 系列可使用对应的 SOC_VERSION。
export SOC_VERSION=ascend910_9391

未设置 CATLASS_SOURCE_DIR 时，构建系统会从 GitCode 自动拉取经过验证的固定 CATLASS 版本。离线构建或 CATLASS 本地开发时可以指定已有源码：

export CATLASS_SOURCE_DIR=/path/to/catlass

编译成功后，可以运行私有 BF16 BasicMatmul smoke 算子验证编译、动态链接和 NPU 执行链路：

python3 - <<'PY'
import torch
import torch_npu
import mindspeed_ops

torch.npu.set_device(0)
torch.manual_seed(42)
a = (torch.rand((128, 128), device="npu") - 0.5).to(torch.bfloat16)
b = (torch.rand((128, 128), device="npu") - 0.5).to(torch.bfloat16)

output = torch.ops.mindspeed_ops._catlass_basic_matmul(a, b)
reference = torch.matmul(a.float(), b.float()).to(torch.bfloat16)
torch.npu.synchronize()

print("output shape:", output.shape)
print("max error:", (output - reference).abs().max().cpu().item())
PY

_catlass_basic_matmul 只用于验证 CATLASS 基础设施，不是正式公开算子接口。目前支持 BF16、二维连续输入， 且 M、N、K 均需按 16 对齐。

CATLASS chunk_loss 的编译、接口和输入约束请参见 chunk_loss 算子说明。

可选：安装 TileLang-Ascend 以使用 TileLang 算子

MindSpeed Ops 默认安装流程不会自动编译或安装 TileLang-Ascend。如需使用 mindspeed_ops.api.tilelang 下的算子，请在已安装 CANN、PyTorch 和 torch_npu 的 Ascend 环境中额外安装 TileLang-Ascend。

从源码构建并安装 TileLang-Ascend wheel：

git clone --recursive https://github.com/tile-ai/tilelang-ascend.git
cd tilelang-ascend

source /usr/local/Ascend/ascend-toolkit/latest/set_env.sh
export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest

./build_wheel_ascend.sh --enable-llvm -v
pip install -U dist/tilelang-*.whl

安装完成后，验证 TileLang-Ascend 和 MindSpeed Ops 均可正常导入：

python -c "import tilelang; print(tilelang.__version__)"
python -c "from mindspeed_ops.api.tilelang.sinkhorn import sinkhorn; print(sinkhorn)"

支持算子清单

算子名称	算子描述	Released	triton	tilelang
FusedCrossEntropyLoss	自定义融合算子	x	√	x
RmsnormWithoutWeight	自定义融合算子	x	√	x
Sinkhorn	Sinkhorn算法	x	√	√
Causal_conv1d	causal_conv1d算子	x	√	x
Wy_fast	wy_fast算子	x	√	x
RmsNormGated	RmsNormGated算子	x	√	x
chunk_kda_bwd_wy_dqkg_fused	KDA chunk backward 融合算子	x	√	x
recompute_w_u_fwd	GDN recompute_w_u_fwd 算子	x	√	x
Sort_chunks_by_idx	MoE chunk重排算子	x	√	x
Clipped_swiglu	GPT-OSS风格SwiGLU-gate融合算子	x	√	x
chunk_bwd_dh	Simple GLA chunk反向dh算子	x	√	x
chunk_bwd_dqkwg	GDN chunk反向dqkwg融合算子	x	√	x
chunk_bwd_dv_local	GDN chunk反向局部dv算子	x	√	x
chunk_fwd_o	GDN chunk前向输出算子	x	√	x
chunk_gated_delta_rule	Gated Delta Rule chunk融合算子	x	√	x
chunk_gated_delta_rule_bwd_dhu	Gated Delta Rule反向dhu算子	x	√	x
chunk_gated_delta_rule_fwd_h	Gated Delta Rule前向隐藏态算子	x	√	x
chunk_kda_fwd_intra	KDA chunk前向intra算子	x	√	x
chunk_local_cumsum	chunk局部cumsum融合算子	x	√	x
chunk_scaled_dot_kkt	chunk缩放点积KKT算子	x	√	x
fused_recurrent_bwd	Simple GLA融合循环反向算子	x	√	x
l2norm	L2归一化算子	x	√	x
mhc_post	MHC后处理算子	x	√	x
mhc_pre_bmm	MHC预处理BMM算子	x	√	x
solve_tril	下三角线性方程求解算子	x	√	x

开发指南

• 新增自定义 Triton 算子指南：介绍 Triton Kernel、Python 封装、UT/ATK 测试以及算子合入所需交付件。
• 新增自定义 AscendC 算子指南：介绍基于 AscendC/ACLNN 实现自定义算子的目录结构、算子开发、测试与合入要求。
• 新增自定义 CATLASS 算子指南：介绍 CATLASS Kernel、CMake/PyTorch 集成、UT/性能测试以及算子合入检查。

分支维护策略

🛠️ MindSpeed Ops 版本分支的维护阶段如下：

状态	时间	说明
计划 🕐	1-3 个月	计划特性
开发 🕔	3 个月	开发特性
维护 🕚	6-12 个月	合入所有已解决的问题并发布版本，针对不同的MindSpeed Ops 版本采取不同的维护策略，常规版本和长期支持版本维护周期分别为6个月和12个月
无维护 🕛	0-3 个月	合入所有已解决的问题，无专职维护人员，无版本发布
生命周期终止（EOL）🚫	N/A	分支不再接受任何修改

算子开发及合入说明

1、算子合入说明请参考ops.md

2、算子开发可参考使用skills工具:

• triton相关： triton算子迁移npu triton在npu上优化
• tilelang相关： tilelang算子迁移npu
• 性能采集skills： 算子性能采集

常见问题

现象	介绍
module 'triton.language' has no attribute 'extract_slice'	问题介绍

免责声明

致MindSpeed Ops使用者

1. MindSpeed Ops提供的所有内容仅供您用于非商业目的。
2. 对于MindSpeed Ops测试用例以及示例文件中所涉及的各模型和数据集,平台仅用于功能测试,华为不提供任何模型权重和数据集,如您使用这些数据进行训练,请您特别注意应遵守对应模型和数据集的License,如您因使用这些模型和数据集而产生侵权纠纷,华为不承担任何责任。
3. 如您在使用MindSpeed Ops过程中,发现任何问题(包括但不限于功能问题、合规问题),请在GitCode提交issue,我们将及时审视并解决。

License声明

Ascend MindSpeed Ops中涉及的算子,如目录下存在License的,以该License为准。如目录下不存在License的,以Apache 2.0许可证许可,对应许可证文本可查阅Ascend MindSpeed Ops根目录。