从源码构建

本页介绍如何构建 Triton-distributed-ascend 及其两个原生依赖项。假设你已经完成 安装,包括 source CANN 的 set_env.sh

构建共分四步,按顺序执行:

  1. 构建 LLVM(可选,仅当没有兼容的预编译 LLVM 时)

  2. 构建 Triton-distributed-ascend

  3. 构建并安装 AscendNPU-IR

  4. 构建并安装 shmem

第 3 步和第 4 步是必需的。AscendNPU-IR 提供将 kernel 降低到 NPU 的编译器后端,shmem 提供设备侧通信运行时。缺少任何一个,软件包虽然能够导入但无法运行分布式 kernel。

备注

以下每一步都需要 CANN 环境。每个 shell 都需要 source 一次: source /usr/local/Ascend/ascend-toolkit/set_env.sh

1. 构建 LLVM(可选)

除非构建过程报告没有兼容的预编译 LLVM,或者你确实需要自定义 LLVM,否则可以跳过本步骤。 构建 LLVM 耗时很长且占用大量磁盘空间。

检出固定的版本

版本与补丁均已固定。使用其他 LLVM commit 将无法正常工作。

git clone --no-checkout https://github.com/llvm/llvm-project.git
cd llvm-project
git checkout f6ded0be897e2878612dd903f7e8bb85448269e5
wget https://raw.githubusercontent.com/triton-lang/triton-ascend/2e69438a0a41a00ab72c7d46b44d97092cf2362c/third_party/ascend/patch/llvm_patch_f6ded0b.patch
git apply llvm_patch_f6ded0b.patch

安装 Clang 与 LLD

建议使用 Clang 和 LLD 15 或更高版本:

apt-get install -y clang-15 lld-15 ccache

配置并构建

设置安装前缀,然后构建:

export LLVM_INSTALL_PREFIX={PATH_TO}/llvm-install
cd {PATH_TO}/llvm-project
mkdir build
cd build
cmake ../llvm \
  -G Ninja \
  -DCMAKE_C_COMPILER=/usr/bin/clang-15 \
  -DCMAKE_CXX_COMPILER=/usr/bin/clang++-15 \
  -DCMAKE_LINKER=/usr/bin/lld-15 \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DLLVM_ENABLE_PROJECTS="mlir;llvm;lld" \
  -DLLVM_TARGETS_TO_BUILD="host;NVPTX;AMDGPU" \
  -DLLVM_ENABLE_LLD=ON \
  -DCMAKE_INSTALL_PREFIX=${LLVM_INSTALL_PREFIX}
ninja install

拷贝测试工具

ninja install 不会安装 FileCheckllvm-lit,但构建过程需要它们:

cp {PATH_TO}/llvm-project/build/bin/FileCheck ${LLVM_INSTALL_PREFIX}/bin/FileCheck
cp {PATH_TO}/llvm-project/build/bin/llvm-lit ${LLVM_INSTALL_PREFIX}/bin/llvm-lit

在下一步构建时通过 LLVM_SYSPATH=${LLVM_INSTALL_PREFIX} 使用此 LLVM。

2. 构建 Triton-distributed-ascend

克隆并初始化子模块

子模块分两个阶段初始化:先初始化外层仓库,然后再初始化 3rdparty/triton-ascend 中的 嵌套子模块。

git clone https://gitcode.com/Ascend/Triton-distributed-ascend.git
cd Triton-distributed-ascend

git submodule update --init --depth=1
cd 3rdparty/triton-ascend
git submodule update --init --depth=1
cd ../../

构建并安装

TRITON_USE_ASCEND=ON \
TRITON_BUILD_WITH_CLANG_LLD=ON \
TRITON_BUILD_PROTON=OFF \
TRITON_BUILD_LITTLE_KERNEL=OFF \
pip install ./python --no-build-isolation

如果在第 1 步中构建了自定义 LLVM,请添加 LLVM_SYSPATH=${LLVM_INSTALL_PREFIX}

备注

必须使用 --no-build-isolation。构建过程需要从你的环境中读取 torch 和 pybind11, 隔离的构建环境中不会有这些依赖。

其他构建命令

如需可编辑安装(在修改 Python 源码时使用):

TRITON_USE_ASCEND=ON \
TRITON_BUILD_WITH_CLANG_LLD=ON \
TRITON_BUILD_PROTON=OFF \
TRITON_BUILD_LITTLE_KERNEL=OFF \
pip install -e ./python --verbose --no-build-isolation

仅生成 wheel 包而不安装:

TRITON_USE_ASCEND=ON \
TRITON_BUILD_WITH_CLANG_LLD=ON \
TRITON_BUILD_PROTON=OFF \
TRITON_BUILD_LITTLE_KERNEL=OFF \
pip wheel ./python --no-build-isolation

构建选项

变量

含义

TRITON_USE_ASCEND=ON

使用 triton-ascend 而非上游 Triton 构建。必需。

TRITON_BUILD_WITH_CLANG_LLD=ON

使用 Clang 和 LLD 进行链接

TRITON_BUILD_PROTON=OFF

禁用 Proton profiler 构建

TRITON_BUILD_LITTLE_KERNEL=OFF

禁用 mega-kernel 构建

TRITON_PLUGIN_DIRS

Triton 插件目录路径

LLVM_SYSPATH

自定义 LLVM 安装路径,来自第 1 步

3. 构建并安装 AscendNPU-IR

AscendNPU-IR 是独立的仓库。请在 Triton-distributed-ascend 目录树之外克隆它。

source /usr/local/Ascend/ascend-toolkit/set_env.sh
git clone https://gitcode.com/Ascend/AscendNPU-IR.git
cd AscendNPU-IR
git submodule update --init --depth=1
mkdir build
./build-tools/build.sh -o ./build -t --build-type Release \
  --apply-patches \
  --bisheng-compiler=$ASCEND_HOME_PATH/bin \
  --build-shmem-template

--bisheng-compiler 指向 CANN 附带的毕昇编译器,因此 ASCEND_HOME_PATH 必须已设置。 如果该变量为空,请重新 source set_env.sh

构建会在 AscendNPU-IR/build/bin 中生成二进制文件,运行时必须将其加入 PATH

export PATH={PATH_TO}/AscendNPU-IR/build/bin:$PATH

请将其加入 shell 配置文件以便在多个会话间持久化。

4. 构建并安装 shmem

shmem 位于仓库的子模块中,提供设备侧通信运行时。

source /usr/local/Ascend/ascend-toolkit/set_env.sh
cd Triton-distributed-ascend/3rdparty/shmem/
bash scripts/build.sh -python_extension
pip install dist/shmem-*.whl

-python_extension 标志会构建 Python 扩展并在 dist/ 中生成 wheel 包。

对于昇腾 Ascend 950 硬件,需添加 SOC 类型:

bash scripts/build.sh -python_extension -soc_type Ascend950

备注

默认构建面向 Ascend910B 后端,覆盖 A2 和 A3 系列。运行 bash scripts/build.sh --help 可查看其他选项,包括 RDMA 支持。

验证构建

检查软件包与通信运行时能否正常导入:

python -c "import triton_dist; print(triton_dist.__file__)"
python -c "import shmem; print('shmem ok')"

然后从仓库根目录运行昇腾单元测试:

pytest python/triton_dist/test/ascend/ -m dist

这些测试直接验证通信原语(barrier、wait/notify、对称内存寻址、put/get),需要至少两张 NPU。

故障排查

ASCEND_HOME_PATH 为空。 你没有在当前 shell 中 source set_env.sh,或者 CANN 安装在 /usr/local/Ascend 以外的位置。调整路径后重新 source。

构建时找不到 LLVM。 你的平台没有兼容的预编译 LLVM。请按第 1 步构建 LLVM 并传入 LLVM_SYSPATH

import triton_dist 成功但 kernel 编译失败。 AscendNPU-IR 二进制文件可能不在 PATH 中。用 which 检查 AscendNPU-IR/build/bin 中的某个二进制文件。

构建时缺少 torch 或 pybind11。 你可能丢掉了 --no-build-isolation,或者没有运行 pip install -r requirements.txt

后续步骤