MindSpeed-Ops

简介

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

版本说明

软件	版本
MindSpeed Ops分支	master
CANN版本	9.0.0
PyTorch	>=2.7.1
triton-ascend	3.2.1
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

可选：安装 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	自定义融合算子	❌	√	❌
RmsnormWithoutWeight	自定义融合算子	❌	√	❌
Sinkhorn	Sinkhorn算法	❌	√	√
Causal_conv1d	causal_conv1d算子	❌	√	❌
Wy_fast	wy_fast算子	❌	√	❌
RmsNormGated	RmsNormGated算子	❌	√	❌
chunk_kda_bwd_wy_dqkg_fused	KDA chunk backward 融合算子	❌	√	❌
recompute_w_u_fwd	GDN recompute_w_u_fwd 算子	❌	√	❌
Sort_chunks_by_idx	MoE chunk重排算子	❌	√	❌
Clipped_swiglu	GPT-OSS风格SwiGLU-gate融合算子	❌	√	❌

分支维护策略

🛠️ 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根目录。