Installation

This page covers the hardware and software prerequisites for Triton-distributed-ascend, and how to prepare your environment before building from source.

Triton-distributed-ascend currently has no prebuilt wheels published, so installation means building from source. Work through this page first, then follow Build from Source.

Hardware requirements

Triton-distributed-ascend targets Ascend products:

  • Atlas A2

  • Atlas A3

  • Ascend 950

Distributed kernels need at least two NPUs visible to the host, since every tutorial and test launches one process per device. A single-device machine is enough to build the project, but not to run the communication examples.

Software requirements

Component

Version

Notes

CANN

9.1.0 or later

Provides the toolkit, the Bisheng compiler, and set_env.sh

Python

3.8 or later

A virtual environment or conda env is recommended

PyTorch

2.0 or later

requirements.txt pins 2.7.1

torch_npu

Matched to your PyTorch and CANN version

requirements.txt pins 2.7.1.post8

CMake

3.19 or later

Required by the C++/MLIR build

GLIBC

2.28 or later

Older distributions cannot load the prebuilt LLVM

libstdc++

3.4.30 or later

Older distributions cannot load the prebuilt LLVM

Note

The PyTorch and torch_npu versions must match. A torch_npu release is built against one specific PyTorch minor version and will fail to import against any other. The pinned pair in requirements.txt (torch 2.7.1 with torch-npu 2.7.1.post8) is the combination this project is tested with.

1. Install CANN

CANN supplies the Ascend toolkit, the Bisheng compiler used to build AscendNPU-IR, and the environment script that every later step sources. Follow the official guide for your OS and install mode:

Install the toolkit package, not just the runtime, because the build needs headers and the compiler binaries.

After installation, source the environment script. The default root is /usr/local/Ascend:

source /usr/local/Ascend/ascend-toolkit/set_env.sh

This sets ASCEND_HOME_PATH, which later build steps rely on. Confirm it is populated:

echo $ASCEND_HOME_PATH

If the variable is empty, the script did not run correctly and the AscendNPU-IR and shmem builds will fail. You need to source this script in every new shell, so consider adding it to your shell profile.

Verify the driver can see your NPUs:

npu-smi info

2. Install PyTorch and torch_npu

Install a PyTorch build and the matching torch_npu plugin for your environment. See the Ascend Extension for PyTorch documentation for the compatibility matrix and installation options.

3. Install Python dependencies

From the repository root:

pip install -r requirements.txt

This installs the pinned torch and torch-npu pair, pybind11 for the Python bindings, and the Python packages CANN itself expects (attrs, cython, numpy, decorator, sympy, cffi, pyyaml, pathlib2, psutil, protobuf, scipy, requests, absl-py).

Note

requirements.txt pulls torch from the CPU wheel index. Device support comes from torch_npu rather than from a device-specific torch build, so this is expected.

Verifying the environment

Before moving on, check that the base stack imports and sees your hardware:

python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__); print(torch.npu.device_count())"

A device count of zero means torch_npu cannot reach the driver. Recheck npu-smi info and confirm you sourced set_env.sh in the current shell.

Next steps