跳转至

Triton 昇腾版（triton-ascend）

约 689 个字 36 行代码 预计阅读时间 3 分钟

Triton 是目前大模型算子开发事实上的主流 DSL（FlashAttention、各类 fused kernel 都用它编写）。triton-ascend 是 triton-lang 官方组织下的昇腾后端，让几乎不用改动的 Triton 代码直接跑在 NPU 上。如果你已经会写 CUDA 上的 Triton，上手成本接近于零。

环境要求

triton-ascend 对版本配套同样敏感（写作时 2026 年 8 月的官方推荐组合，以官方 README 为准）：

组件 版本
CANN 9.1.0
Python 3.9–3.11（推荐 3.11）
torch_npu 2.7.1.post8
硬件 Atlas A2/A3/950 系列，建议单卡 ≥ 32 GB

安装

pip 安装（推荐），从昇腾的 PyPI 镜像源拉取：

pip install triton-ascend --extra-index-url=https://mirrors.huaweicloud.com/ascend/repos/pypi

源码安装适合想跟进开发的场景：需要 clang-15/lld-15 等编译依赖，然后

git clone https://github.com/triton-lang/triton-ascend.git
cd triton-ascend
pip install -e .

无论哪种方式，都要先激活 CANN 环境：

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

容器更省心

官方文档推荐直接使用预装 CANN 的开发镜像（如 quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.11），再在容器内执行上述安装，参见环境搭建。

小试牛刀：向量加法

下面改写自官方教程 01-vector-add.py（略有简化）。与 NVIDIA Triton 的写法逐行对照，你会发现唯一的区别就是张量创建在 npu 上：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import torch
import torch_npu

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)


def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output


x = torch.rand(98432, device='npu')
y = torch.rand(98432, device='npu')
torch.testing.assert_close(add(x, y), x + y)
print("======Vector Add Test Passed!======")

能做什么、和 CUDA Triton 差在哪

  • 能做什么：向量算子、矩阵乘（Cube）、融合算子（如激活+归一化）都能写，tl.* 的高级原语（autotune、reduce、dot 等）可用；
  • 自动切分：昇腾片上存储有限，triton-ascend 编译器会自动对数据做 tiling 与多级缓冲管理，省去了 Ascend C 手写搬运流水的功夫，代价是对底层流水线的控制粒度略低；
  • autotune：@triton.autotune 按同样语法使用，会在 NPU 上自动搜索 block 尺寸等配置；
  • 官方文档站 triton-ascend.readthedocs.io 提供 Vector / Cube / 融合算子三类教程与典型算子样例（向量加法、矩阵乘、注意力等），建议按顺序过一遍。

常见坑

如果报 ModuleNotFoundError: No module named 'triton._C.libtriton.ascend'，多半是后装的 NVIDIA 版 Triton 覆盖了 triton-ascend 的目录，卸载后重装 triton-ascend 即可。同一环境里尽量不要混装两个 Triton。

学习&拓展：

  1. 跑通向量加法样例，并故意把向量长度改成非整数个 block，观察 mask 的作用。
  2. 用 @triton.autotune 对矩阵乘 kernel 的 BLOCK_M/BLOCK_N/BLOCK_K 做自动调优，记录最优配置。
  3. 对照官方教程用 tl.dot 写一个矩阵乘，体会它何时会走 Cube 单元。
  4. 阅读官方“融合算子开发”教程，实现一个 linear + GELU 融合算子，并与“先 linear 后 GELU”两步执行对比耗时。
  5. 思考：同样是 FlashAttention，Triton 版本与 Ascend C 版本各自把复杂度藏在了哪里？

一些链接

triton-ascend 官方仓库（GitHub） / Gitee 镜像

中文安装指南

官方文档站（Vector/Cube/融合算子教程）