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

MindSpeed Ops

MindSpeed Ops is a custom operator library for Ascend NPU training workloads. It focuses on fused and hotspot operators in large-model training and provides a unified PyTorch-style Python API (mindspeed_ops.api.*), helping training workloads gain performance benefits on the Ascend platform at a low integration cost.

MindSpeed Ops provides four types of operator implementations: Triton operators built on Triton-Ascend, provided with the default installation and covering commonly used fused operators in training workloads; TileLang operators built on TileLang-Ascend, targeting specific scenarios such as vector core and requiring additional installation; ACLNN operators built on CANN AscendC, which are compiled into a dynamic library for the chip model (SOC_VERSION) during installation and registered with torch.ops.mindspeed_ops; and Ascend C operators based on CATLASS, which are not compiled by default and require MINDSPEED_BUILD_CATLASS to be enabled; currently, they are supported only on Atlas A2/A3 training products. Operator implementations are provided for two architectures: arch32 (Atlas A2/A3 training products) and arch35 (Ascend 950PR&950DT products). The chip model is detected at runtime and requests are dispatched to the corresponding implementation, so the same calling code runs on different chip platforms.

All operator interfaces follow the native PyTorch calling conventions (torch.autograd.Function or regular functions) and can be embedded directly into existing training scripts and model code. Each operator comes with accuracy-alignment unit tests (tests/unit_tests/) and accuracy, performance, and memory tests through ATK (tests/atk_tests/).

The MindSpeed Ops documentation includes the following core content:

Project Introduction: Introduction - framework positioning, architecture, and features.
Environment Installation Guide: Installation Guide - dependency compatibility and installation configuration.
Quick Start: Quick Start - a step-by-step guide from installing and verifying the operator library to invoking operators for a computation task.
Supported Operators: Supported Operators - supported operators.
Operator Development Guide:
Docker Image: MindSpeed Ops Docker Image Overview

Release Notes

See Release Notes for the latest and historical MindSpeed Ops releases, including the software compatibility matrix, version compatibility notes, and the update notes for the current version.

Installation

For detailed installation steps and environment configuration, see the MindSpeed Ops Installation Guide.

Supported Operators

Operator	Description	Released	triton	tilelang	catlass
`FusedCrossEntropyLoss`	Custom fused operator	x	√	x	x
`RmsnormWithoutWeight`	Custom fused operator	x	√	x	x
`Sinkhorn`	Sinkhorn algorithm	x	√	√	x
`Causal_conv1d`	causal_conv1d operator	x	√	x	x
`Wy_fast`	wy_fast operator	x	√	x	x
`RmsNormGated`	RmsNormGated operator	x	√	x	x
`chunk_kda_bwd_wy_dqkg_fused`	KDA chunk backward fused operator	x	√	x	x
`recompute_w_u_fwd`	GDN recompute_w_u_fwd operator	x	√	x	x
`Sort_chunks_by_idx`	MoE chunk reordering operator	x	√	x	x
`Clipped_swiglu`	GPT-OSS-style SwiGLU-gate fused operator	x	√	x	x
`chunk_bwd_dqkwg`	GDN chunk backward dqkwg fused operator	x	√	x	x
`chunk_bwd_dv_local`	GDN chunk backward local dv operator	x	√	x	x
`chunk_fwd_o`	GDN chunk forward output operator	x	√	x	x
`chunk_gated_delta_rule`	Gated Delta Rule chunk fused operator	x	√	x	x
`chunk_gated_delta_rule_bwd_dhu`	Gated Delta Rule backward dhu operator	x	√	x	x
`chunk_gated_delta_rule_fwd_h`	Gated Delta Rule forward hidden-state operator	x	√	x	x
`chunk_kda_fwd_intra`	KDA chunk forward intra operator	x	√	x	x
`chunk_local_cumsum`	Chunk local cumsum fused operator	x	√	x	x
`chunk_scaled_dot_kkt`	Chunk scaled dot-product KKT operator	x	√	x	x
`l2norm`	L2 normalization operator	x	√	x	x
`mhc_post`	MHC post-processing operator	x	√	x	x
`mhc_pre_bmm`	MHC pre-processing BMM operator	x	√	x	x
`solve_tril`	Lower-triangular linear equation solver	x	√	x	x
`inplace_rope`	DeepSeek V4 interleaved RoPE operator	x	√	x	x
`chunk_loss`	Chunked language-model cross-entropy operator (including backward gradients for hidden states and vocabulary weights)	x	x	x	√

Operator Development and Integration

To add a new operator and integrate it into this repository, see the Operator Integration Guide and the following operator development guides.

Custom Triton Operator Development Guide: describes the Triton Kernel, Python wrapper, UT/ATK tests, and the deliverables required for operator integration.
Custom AscendC Operator Development Guide: describes the directory structure, operator development, testing, and integration requirements for custom operators implemented with AscendC/ACLNN.
Custom CATLASS Operator Development Guide: describes the CATLASS Kernel, CMake/PyTorch integration, UT/performance tests, and operator integration checks.
Custom ATB Operator Development Guide: describes the ATB Operation, JIT Builder, Python API, packaging, UT/ATK tests, and operator integration checks.
Custom TileLang Operator Development Guide: describes implementing a TileLang Kernel, adding a Python calling interface, writing accuracy UTs, and adding ATK tests.

The following skills tools can be referenced during operator development:

Triton-related: triton operator migration to NPU triton optimization on NPU
TileLang-related: tilelang operator migration to NPU
Performance profiling skills: operator performance profiling

Version Maintenance Policy

🛠️ MindSpeed Ops versions go through the following five maintenance phases.

Phase	Time	Description
Planning 🕐	1-3 months	Plans features
Development 🕔	3 months	Develops features
Maintenance 🕚	6-12 months	Merges all resolved issues and releases versions. Different MindSpeed Ops versions follow different maintenance policies; the maintenance periods for regular and long-term support versions are 6 and 12 months, respectively.
No Maintenance 🕛	0-3 months	Merges all resolved issues. No dedicated maintainer is assigned, and no versions are released.
End of Life (EOL)	N/A	The branch no longer accepts any changes.

FAQ

For common questions about basic repository usage, see the MindSpeed Ops FAQ. If the FAQ does not cover your question, search the repository's issues list for similar problems or submit a new issue.

Security Statement

MindSpeed Ops Security Statement.

Disclaimer

To MindSpeed Ops Users

All content provided by MindSpeed Ops is for non-commercial use only.
The platform is used only for functional testing of the models and datasets involved in the MindSpeed Ops test cases and example files. Huawei does not provide any model weights or datasets. If you use these data for training, you must comply with the licenses of the corresponding models and datasets. Huawei assumes no responsibility for any infringement disputes arising from your use of these models and datasets.
If you encounter any problems while using MindSpeed Ops (including but not limited to functional and compliance issues), submit an issue on GitCode. We will review and resolve it promptly.

License Statement

For the operators involved in Ascend MindSpeed Ops, if a LICENSE file exists in the directory, that license prevails. If no LICENSE file exists, the operators are licensed under the Apache License 2.0, whose text can be found in the root directory of Ascend MindSpeed Ops.

Contribution Statement

To report issues or contribute code to MindSpeed Ops, see the Contribution Guide.

Acknowledgments

Thank you to everyone in the community for every PR. Contributions to MindSpeed Ops are welcome.

Introduction