MindSpeed Ops

简介

MindSpeed Ops是面向昇腾NPU的训练业务自定义算子库，聚焦大模型训练中的融合算子与热点算子，提供PyTorch风格的统一Python API（mindspeed_ops.api.*），帮助训练业务以较低的接入成本获得昇腾平台上的性能收益。

MindSpeed Ops提供四类算子实现：基于Triton-Ascend的Triton算子，随默认安装提供，覆盖训练业务高频融合算子；基于TileLang-Ascend的TileLang算子，面向vector core等特定场景，需额外安装后使用；基于CANN AscendC的ACLNN算子，安装时按芯片型号（SOC_VERSION）编译为动态库并注册到torch.ops.mindspeed_ops；基于CATLASS的Ascend C算子，默认不编译，需开启MINDSPEED_BUILD_CATLASS编译，当前仅支持Atlas A2/A3训练系列产品。算子实现区分arch32（Atlas A2/A3训练系列产品）与arch35（Ascend 950PR&950DT系列产品）两套架构，运行时自动检测芯片型号并分发到对应实现，同一份调用代码可跨芯片平台运行。

所有算子接口与PyTorch原生调用方式一致（torch.autograd.Function或普通函数），可直接嵌入现有训练脚本与模型代码；每个算子配套精度对齐单元测试（tests/unit_tests/），并通过ATK完成精度、性能与内存测试（tests/atk_tests/）。

详细文档请参见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算子当前仅支持Atlas A2训练系列产品与Atlas A3训练系列产品，Ascend 950PR&950DT系列产品暂不支持。 在已安装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
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	catlass
FusedCrossEntropyLoss	自定义融合算子	x	√	x	x
RmsnormWithoutWeight	自定义融合算子	x	√	x	x
Sinkhorn	Sinkhorn算法	x	√	√	x
Causal_conv1d	causal_conv1d算子	x	√	x	x
Wy_fast	wy_fast算子	x	√	x	x
RmsNormGated	RmsNormGated算子	x	√	x	x
chunk_kda_bwd_wy_dqkg_fused	KDA chunk backward融合算子	x	√	x	x
recompute_w_u_fwd	GDN recompute_w_u_fwd算子	x	√	x	x
Sort_chunks_by_idx	MoE chunk重排算子	x	√	x	x
Clipped_swiglu	GPT-OSS风格SwiGLU-gate融合算子	x	√	x	x
chunk_bwd_dqkwg	GDN chunk反向dqkwg融合算子	x	√	x	x
chunk_bwd_dv_local	GDN chunk反向局部dv算子	x	√	x	x
chunk_fwd_o	GDN chunk前向输出算子	x	√	x	x
chunk_gated_delta_rule	Gated Delta Rule chunk融合算子	x	√	x	x
chunk_gated_delta_rule_bwd_dhu	Gated Delta Rule反向dhu算子	x	√	x	x
chunk_gated_delta_rule_fwd_h	Gated Delta Rule前向隐藏态算子	x	√	x	x
chunk_kda_fwd_intra	KDA chunk前向intra算子	x	√	x	x
chunk_local_cumsum	chunk局部cumsum融合算子	x	√	x	x
chunk_scaled_dot_kkt	chunk缩放点积KKT算子	x	√	x	x
l2norm	L2归一化算子	x	√	x	x
mhc_post	MHC后处理算子	x	√	x	x
mhc_pre_bmm	MHC预处理BMM算子	x	√	x	x
solve_tril	下三角线性方程求解算子	x	√	x	x
chunk_loss	分块语言模型交叉熵算子（含hidden states和词表权重的反向梯度）	x	x	x	√

开发指南

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

分支维护策略

🛠️ 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： 算子性能采集

贡献声明

如果您希望向MindSpeed Ops报告问题和贡献代码，具体请参见贡献指南。

常见问题

现象	介绍
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根目录。