Triton-Ascend is an optimized version of Triton that adapts to Huawei Ascend chips. It provides efficient automatic optimization of kernel functions, operator compilation, and deployment capabilities, and supports products such as Ascend Atlas A2/A3. While being compatible with the core syntax of Triton, Ascend is optimized for features of Ascend NPUs, including automatic parsing of kernel function parameters, memory access logic optimization, and security deployment mechanism optimization.
Complete online documents and network materials are provided, covering environment setup, operator development, optimization practices, and FAQ, to help you get started quickly. For details, see the online documents.
Supported OS: Linux (AArch64/x86_64)
Supported Ascend products: Atlas A2/A3 series
Minimum hardware configuration: single-device 32 GB graphics memory (recommended)
Python (Python 3.9 to Python 3.13,Tips:Python 3.9 does not support AArch64), CANN_TOOLKIT, CANN_OPS, requirements.txt, and requirements_dev.txt
For details about the CANN installation and configuration script, see CANN installation description. Developers can select the CANN version, product series, CPU architecture, operating system, and installation method to find the corresponding installation commands.
You need to select the required version (9.0.0 is recommended) based on the Ascend card model you use. The CANN installation takes about 5 to 10 minutes. Wait until the installation is complete.
- Note: If the user does not specify an installation path, the software will be installed to the default path as follows.
For root user:
/usr/local/AscendFor non-root user:${HOME}/Ascend, where${HOME}is the home directory of the current user.
The environment variable configuration above takes effect only in the current terminal session.
Users can add the command source ${HOME}/Ascend/ascend-toolkit/set_env.sh to an environment variable configuration file (such as .bashrc) as needed.
You can run the following command to install the requirements:
pip install -r requirements.txt -r requirements_dev.txtYou can set up the Triton-Ascend environment by referring to section "Preparing the Environment" in Installation Guide.
You can install the latest stable version package using the CLI.
pip install triton-ascend-
Note: Starting from version 3.5, Triton-Ascend mitigates the installation overwrite issue by declaring Triton as an installation dependency. When Triton-Ascend is installed, community Triton is installed first, and then Triton-Ascend overwrites the shared package directory. This helps avoid a later Triton reinstallation overwriting Triton-Ascend when other software packages that depend on Triton are installed. Different community Triton package versions are used for x86 and arm because arm installation packages are only available in the community starting from Triton 3.5: x86 depends on
triton==3.2.0, and arm depends ontriton==3.5.0. -
Note 1: This solution mitigates the installation overwrite issue, but it does not completely eliminate the conflict caused by community Triton and Triton-Ascend sharing the same top-level
tritonpackage directory. If a later installation explicitly reinstalls or upgrades community Triton, the installed Triton-Ascend may still be affected. In this case, uninstall both community Triton and Triton-Ascend first, and then reinstall Triton-Ascend.
You can also download the nightly package from the download link and install it locally.
- Note 2: If you download the nightly package for installation, select the Python version and architecture (AArch64/x86_64) of your server when selecting the Triton-Ascend package.
- Note 3: The nightly package is built every day. Developers submit MRs frequently. Note that if the package does not pass the stable test, function bugs may exist.
We provide a Dockerfile to help you build a Docker environment image. The build uses pre-built CANN images from quay.io/ascend/cann as the base, which significantly speeds up the build process by skipping the CANN installation step.
You need to specify the CANN_BASE_IMAGE build arg to select the appropriate CANN base image for your machine. Available CANN base image tags can be found at quay.io/ascend/cann.
| CANN Version | Chip Type | Python | Image Tag |
|---|---|---|---|
| 8.5.0 | A2 |
3.10 | 8.5.0-910b-ubuntu22.04-py3.10 |
| 8.5.0 | A3 |
3.10 | 8.5.0-a3-ubuntu22.04-py3.10 |
| 8.5.0 | A2 |
3.11 | 8.5.0-910b-ubuntu22.04-py3.11 |
| 8.5.0 | A3 |
3.11 | 8.5.0-a3-ubuntu22.04-py3.11 |
| 9.0.0-beta.2 | A2 |
3.10 | 9.0.0-beta.2-910b-ubuntu22.04-py3.10 |
| 9.0.0-beta.2 | A3 |
3.10 | 9.0.0-beta.2-a3-ubuntu22.04-py3.10 |
| 9.0.0-beta.2 | A2 |
3.11 | 9.0.0-beta.2-910b-ubuntu22.04-py3.11 |
| 9.0.0-beta.2 | A3 |
3.11 | 9.0.0-beta.2-a3-ubuntu22.04-py3.11 |
You can check the NPU model on your system using the npu-smi command.
For the machines corresponding to different chip types, refer to the table below:
| Option No. | Chip Type | Corresponding Server/Product Series | Typical Server Model |
|---|---|---|---|
| 1 | A3 |
Atlas A3 Training Series | Atlas 900 A3 SuperPoD |
| 2 | A2 |
Atlas A2 Training Series | Atlas 800T A2 |
git clone https://github.com/triton-lang/triton-ascend.git && cd triton-ascend
docker build \
--build-arg CANN_BASE_IMAGE=quay.io/ascend/cann:8.5.0-a3-ubuntu22.04-py3.10 \
-t triton-ascend-image:latest -f ./docker/Dockerfile .To start a container from this image, you can use the following command as a reference:
docker run -u 0 -dit --shm-size=512g --name=triton-ascend_container --net=host --privileged \
--security-opt seccomp=unconfined \
--device=/dev/davinci0 \
--device=/dev/davinci1 \
--device=/dev/davinci2 \
--device=/dev/davinci3 \
--device=/dev/davinci4 \
--device=/dev/davinci5 \
--device=/dev/davinci6 \
--device=/dev/davinci7 \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /home:/home \
-v /etc/ascend_install.info:/etc/ascend_install.info \
triton-ascend-image:latest \
/bin/bash
# Enter the container
docker exec -u root -it triton-ascend_container /bin/bashRun the example: 01-vector-add.py
# Set CANN environment variables (using the default root installation path `/usr/local/Ascend` as an example)
source /usr/local/Ascend/ascend-toolkit/set_env.sh
# Clone the triton-ascend repository and examples (optional; required for running examples if not installed from source)
git clone https://github.com/triton-lang/triton-ascend.git
# Run the tutorials example:
python3 ./triton-ascend/third_party/ascend/tutorials/01-vector-add.pyOutput similar to the following indicates that your environment is correctly configured.
tensor([0.8329, 1.0024, 1.3639, ..., 1.0796, 1.0406, 1.5811], device='npu:0')
tensor([0.8329, 1.0024, 1.3639, ..., 1.0796, 1.0406, 1.5811], device='npu:0')
The maximum difference between torch and triton is 0.0Triton-Ascend maintains full syntax compatibility with the community Triton, requiring only changes to tensor device declarations and a few torch.cuda.* APIs to run existing GPU examples on Ascend NPU. Below is a complete migration process demonstrated through a typical vector addition test.
The GPU version example file test_add.py is as follows:
import pytest
import torch
from torch.testing import assert_close
import triton
import triton.language as tl
@triton.jit
def add_kernel(
x_ptr,
y_ptr,
output_ptr,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
tl.store(output_ptr + offsets, x + y, mask=mask)
@pytest.mark.parametrize('SIZE,BLOCK_SIZE', [(98432, 1024)])
def test_add(SIZE, BLOCK_SIZE):
device_id = torch.cuda.current_device()
device = torch.device('cuda', device_id)
x = torch.randn(SIZE, device='cuda', dtype=torch.float32)
y = torch.randn(SIZE, device='cuda', dtype=torch.float32)
output_cpu = torch.empty(SIZE, dtype=torch.float32)
output = output_cpu.cuda()
def grid(meta):
return (triton.cdiv(SIZE, meta['BLOCK_SIZE']),)
add_kernel[grid](x, y, output, SIZE, BLOCK_SIZE=BLOCK_SIZE)
torch.cuda.synchronize()
output_torch = x + y
assert_close(output, output_torch, rtol=1e-3, atol=1e-3)To migrate, simply replace the GPU-related APIs with their corresponding NPU versions. The mapping is as follows:
| GPU Syntax | NPU Syntax |
|---|---|
device='cuda' |
device='npu' |
tensor.cuda() |
tensor.npu() |
torch.cuda.current_device() |
torch.npu.current_device() |
torch.cuda.synchronize() |
torch.npu.synchronize() |
The kernel functions annotated with @triton.jit use the Triton generic language and generally do not require special modifications. The Launch grid calling method is also completely consistent with the GPU version.
The core changes are shown in diff format:
import pytest
import torch
from torch.testing import assert_close
import triton
import triton.language as tl
# ...（kernel code remains unchanged）...
@pytest.mark.parametrize('SIZE,BLOCK_SIZE', [(98432, 1024)])
def test_add(SIZE, BLOCK_SIZE):
- device_id = torch.cuda.current_device()
+ device_id = torch.npu.current_device()
- x = torch.randn(SIZE, device='cuda', dtype=torch.float32)
- y = torch.randn(SIZE, device='cuda', dtype=torch.float32)
+ x = torch.randn(SIZE, device='npu', dtype=torch.float32)
+ y = torch.randn(SIZE, device='npu', dtype=torch.float32)
output_cpu = torch.empty(SIZE, dtype=torch.float32)
- output = output_cpu.cuda()
+ output = output_cpu.npu()
add_kernel[grid](x, y, output, SIZE, BLOCK_SIZE=BLOCK_SIZE)
- torch.cuda.synchronize()
+ torch.npu.synchronize()
output_torch = x + y
assert_close(output, output_torch, rtol=1e-3, atol=1e-3)
```bash
After making the modifications, run the test case with `pytest`. Successful execution indicates a successful migration.pytest test_add.py
If the `pytest` component is not installed, you can install it using `pip install pytest`.