Triton学习 · Part 1 · Hello, world!

yLDeveloper2026-05-12 22:06

不重要的碎碎念：

其实早有预谋想要接触AI Infra相关的内容，毕竟虽然属于深度学习领域，但又因为和底层硬件打交道所以没有这么热门，而且他的长辈 - 高性能计算 又是和我专业比较密切的部分，而且相较于一时的业务岗（例如Agent开发）之类的不可替代性（就本文编写时间 2026年4月30日）还比较强。

但是苦于还没有系统学习计组、编译原理等等，很难深入AI Infra的其他领域，于是目前只能先学习与软件关联性更强的算子开发，学习中慢慢加深对计算机的认知。

至于学习动机？华为举行了一个CANN算子开发大赛。因为首届举行所以参赛地区少，人数也少，似乎比较容易混个【谢谢参与】。其使用的 Ascend C 语言只支持华为自己的 NPU，显然我只能发动纯粹的 Vibing 之术（注：其实好像可以申请算力）。即使是我提供的思路，终究交上的代码不是自己亲手敲下的、自己也运行不了；产出还没热乎，就只能交给测评器，而代码到底在人家那里经历了什么也毫不清楚，竟有种苦主之感。这便郁闷了。然而审查着代码发现竟一点看不懂也便更郁闷了。真该狠狠学学算子开发了。

至于为何选择 Triton ，一是难度较低（与CUDA相比），二则是不被老黄霸占，甚至开始支持国产的"A卡"了，在目前的全球格局下的确是一支潜力股。

安装配置

环境：

  • 系统：Manjaro
  • GPU：NVIDIA GeForce MX450
  • CUDA 最高支持版本：13.1

这里没有使用NVIDIA Triton Docker容器，而是直接使用了

bash 复制代码 
source .venv/bin/activate
pip install torch triton -i https://pypi.tuna.tsinghua.edu.cn/simple

Hello, world !

来看一下 Triton 版的 Helloworld：

python 复制代码 
import torch
import triton
import triton.language as tl

@triton.jit
def vector_add_kernel(a_ptr, b_ptr, c_ptr):
    offsets = tl.arange(0, 16)
    a = tl.load(a_ptr + offsets)
    b = tl.load(b_ptr + offsets)
    c = a + b
    tl.store(c_ptr + offsets, c)

def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
    grid = (1,)
    vector_add_kernel[grid](a, b, c)

if __name__ == "__main__":
    N = 16
    a = torch.randn(N, device='cuda')
    b = torch.randn(N, device='cuda')
    triton_output = torch.empty_like(a)
    solve(a, b , triton_output, N)
    print("Answer:", triton_output)

实现了一个简单的，固定向量维度（为16）的向量加法算子。

看不懂？没关系。

下面我们进行分析：

分析

先来复习一下向量加法：

若

\[\text{dim}(\vec{a})=\text{dim}(\vec{b})=n \]

那么有

\[\begin{split} &\vec{a}+\vec{b}\\ =&(a_1,a_2,a_3,\ldots,a_n)+(b_1,b_2,b_3,\ldots,b_n)\\ =&(a_1+b_1,a_2+b_2,a_3+b_3,\ldots,a_n+b_n)\\ \end{split}\]

如果用我们常规的思维（CPU思维）求解，肯定是这么计算这样的加法的（伪代码）：

python 复制代码 
for i in range(0, n):
	c[i] = a[i] + b[i]

这样计算，向量如果有 \(n\) 维度，那么 CPU 就要计算 \(n\) 次。

为了加速这样的运算，可以考虑考虑用多核一起计算并且汇总到 c 上。

这里专门这样并行计算的 GPU 就派上用场了。

假如 GPU有 \(k\) 个线程，那么我们可以这样命令前 \(n\) 个线程：

复制代码 
第一个线程：
计算 a[0] + b[0]

第二个线程：
计算 a[1] + b[1]

第三个线程：
计算 a[2] + b[2]

...

第n个线程：
计算 a[n - 1] + b[n - 1]

放在 CUDA 里面，我们可以为单个线程写脚本，像

cpp 复制代码 
int i = threadIdx.x; 
// 线程自己知道自己的编号 
c[i] = a[i] + b[i]; 
// 只算 1 个标量

但是很显然，这会非常繁琐。成为这样的微操大师，对于暂时没有想法产出凹到机制、榨干硬件的算子的我们没有必要。

Triton所操纵的，不是线程级别的代码，而是向量级的代码。

现在我们回到代码本身。显然代码由3部分构成：

python 复制代码 
import torch
import triton
import triton.language as tl


# 第一部分：算子内核本体
@triton.jit
def vector_add_kernel(a_ptr, b_ptr, c_ptr):
    offsets = tl.arange(0, 16)
    a = tl.load(a_ptr + offsets)
    b = tl.load(b_ptr + offsets)
    c = a + b
    tl.store(c_ptr + offsets, c)

# 第二部分：调用算子的函数
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
    grid = (1,)
    vector_add_kernel[grid](a, b, c)

# 实际开发场景中算子的应用
if __name__ == "__main__":
    N = 16
    a = torch.randn(N, device='cuda')
    b = torch.randn(N, device='cuda')
    triton_output = torch.empty_like(a)
    solve(a, b , triton_output, N)
    print("Answer:", triton_output)
算子内核章

暂且先忽略这个函数是如何运行的，先看看他的思路：

首先可知参数是3个指针，指向加向量 \(\vec{a}\) 、\(\vec{b}\) 和结果输出向量 \(\vec{c}\) 。

然后我们使用 arange 获取了 \(\vec{a}\) 、\(\vec{b}\) 维度的索引（一维到十六维），同样储存在张量里面。

易知，\(\vec{a}\) 、\(\vec{b}\) 的指针地址加上第 \(n\) 维元素的偏移等于指向该元素地址的指针，

显然， offsets 是一个向量（真的么？），（此处可为了理解类比为数组），而一个整数加上一个向量等于一个新向量（原向量的每一位加上该整数）。

所以我们使用了 tl.load 加载了 \(\vec{a}\) 、\(\vec{b}\) 每个元素 （真的吗？）的数值。

紧接着我们进行加法运算，然后将结果储存在 \(c\) 中的每一位上 （同样，真的吗？）

如果第一眼看上去，事先不了解 Triton，似乎可以得出一个结论：

工作流程：main 函数调用 solve -> solve 调用 vector_add_kernel -> vector_add_kernel 调用了向量的加法 （？？？）

那么也就是说程序最后还是执行了不知道用哪个数学库里面的向量加法？还是我们这个"向量加法"算子进入了一个无限循环？

这是一种 CPU 编程思维，如果你尝试在 vector_add_kernel 中加入一个调试信息，就会发现实际上这个调试信息输出了近百次！

其实 vector_add_kernel 并不能被定义为一个普通的函数。它被 @triton.jit 修饰，是一个内核。

这个内核被并不会直接被 Python 的解释器调用，而是会在 Triton 的 JIT（即时编译）下被转为对每个GPU中线程的具体命令。（详情见：Triton学习 · 番外篇 · 编译流程）也就是说你内核写的算子代码（向量级的代码）并不会被GPU中的某个线程所看到。单个特定的线程所看到的，是一个针对数组中某个指定索引元素的操作。如下：

你写的（给编译器看的）

python 复制代码 
offsets = tl.arange(0, 16)       # 向量 [0,1,2,...,15]
a = tl.load(a_ptr + offsets)     # 向量加载
c = a + b                        # 向量运算
tl.store(c_ptr + offsets, c)     # 向量存储

编译器翻译后，每个线程看到的（举例、伪代码）

c 复制代码 
// Thread 0 看到的：
a = load(a_ptr + 0)
c = a + b
store(c_ptr + 0, c)

// Thread 1 看到的：
a = load(a_ptr + 1)
c = a + b
store(c_ptr + 1, c)

// Thread 5 看到的：
a = load(a_ptr + 5)
c = a + b
store(c_ptr + 5, c)

// 并不代表真实的索引分配

这样，每个线程所进行的就是纯正的标量（数值）运算而非向量运算。

offsets 这个向量根本不存在于最终的线程代码中，它被拆解成了每个线程各自的具体数值。

更准确地说，实际的 PTX 大致是这样的：

复制代码 
mov.u32 %r1, %tid.x           // 拿到自己的线程编号
add.u32 %r2, a_ptr, %r1       // a_ptr + 线程编号（就是你的 offset）
ld.global.f32 %f1, [%r2]      // 加载 a
ld.global.f32 %f2, [%r3]      // 加载 b
add.f32 %f3, %f1, %f2         // 计算 a + b
st.global.f32 [%r4], %f3      // 存储 c

其中%tid.x（线程编号）直接替代了原本的 offsets。

调用算子de函数章

现在终于大致理解了算子内核工作是什么原理了，来看看如何调用这个算子吧：

python 复制代码 
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
    grid = (1,)
    vector_add_kernel[grid](a, b, c)

你可能会感到困惑：

这第二句，什么叫一个函数 的数组 的元组 索引？

其实这是 Python 的语法糖，因为Triton 重载了 [] 运算符。

如果展开，那就是：

python 复制代码 
vector_add_kernel[grid](a, b, c)
# 等价于
vector_add_kernel.__getitem__(grid).__call__(a, b, c)

# 第二步：kernel[grid] 返回一个「可调用的启动器」
launcher = vector_add_kernel[grid]   # 类似 __getitem__

# 第三步：调用启动器，传入参数
launcher(a, b, c)

然后我们就通过这个 launcher 去启动这个 GPU 的内核算子。

那么这个 grid 元组又是什么？

为了更好的理解，我们最好先从 CUDA 说起

CUDA 启动 kernel 的语法：

cpp 复制代码 
kernel<<<grid, block>>>(args);
//grid:  几个block
//block: 每个block几个线程

Triton 对应的语法：

python 复制代码 
kernel[grid](args)
# grid:  几个 program（block）

Triton 没有 block 参数，这也是它与 CUDA 的区分之一，因为 block 内部的线程分配由编译器决定。你只需要告诉它启动几个 program。

而 grid 就是一个普通的元组。

python 复制代码 
grid = (4,)         # 启动 4 个 program，1维
grid = (4, 2)       # 启动 4×2=8 个 program，2维
grid = (4, 2, 3)    # 启动 4×2×3=24 个 program，3维

而每个 program 可以通过 tl.program_id(axis) 知道自己是第几号：

python 复制代码 
pid_x = tl.program_id(0)   # 第 0 维的编号
pid_y = tl.program_id(1)   # 第 1 维的编号

我们现在大致理解这个 Triton 版的 Hello, world! 程序了。那么这个程序都有什么问题呢？

开始改进

首先，加向量 \(\vec{a}\) 、\(\vec{b}\) 的元素个数，肯定不能被固定下来。

假如维度为 \(15\)，即主程序的 N=15，是否只要改为 offsets = tl.arange(0, 15) 便能完美解决？你很快便会发现报错了（ValueError: arange's range must be a power of 2）！

注：Triton 的一个 program 大致对应 CUDA 中的一个 thread block/CTA。CTA 会被分配到某个 SM 上执行，而 SM 内部通常以 warp 为单位调度指令。

这并不是因为你令N=15，毕竟向量的维度也不是你能控制的。问题出在传入 arange 的 block 大小有问题。这是因为在常见写法中，tl.arange(0, BLOCK_SIZE) 的范围大小通常需要是 \(2\) 的幂。因此我们通常把每个 program 的 tile 大小 BLOCK_SIZE 设成 \(2\) 的幂，再用 mask 处理真实长度不是 \(2\) 的幂的情况。为什么呢？

我们把一个 program/block 覆盖的逻辑 tile 大小称为 BLOCK_SIZE

  1. NVIDIA GPU 中，SM 内部通常以 warp 为单位调度和执行指令，\(1\) 个 warp 包含 \(32\) 个线程。由于 \(32\) 是 \(2\) 的幂，使用 \(2\) 的幂作为 BLOCK_SIZE 更容易进行规则的线程映射和内存访问优化。
  2. 二的幂方便编译器优化，等等

因此，在这种使用 tl.arange(0, BLOCK_SIZE) 构造 block tensor 的常见写法中，BLOCK_SIZE 通常应设为 \(2\) 的幂。

这里，如果你直接令 BLOCK_SIZE=15 ，并写成 tl.arange(0, BLOCK_SIZE) ，那么 tl.arange(0, 15) 会因为范围大小不是 \(2\) 的幂而报错。

注意，这并不代表 Triton 不能处理长度为 \(15\) 的向量，而是不能这样构造 block tensor。

显然，不能因为这样就放弃那些长度不为 \(2\) 的幂的向量。我们不妨思考一下解决方法。

用掩码处理不足一个 block 的部分

不改动 BLOCK_SIZE ，即依旧使用 tl.arange(0, 16)。但是我们使用 mask 跳过第 \(16\) 个元素，只处理前 \(15\) 个位置，即只处理 offset < 15 的位置（毕竟 offset 一定 \(\textless 16\) ）。

得出代码：

python 复制代码 
@triton.jit
def vector_add_kernel(a_ptr, b_ptr, c_ptr):
    offsets = tl.arange(0, 16)
    mask = offsets < 15
    a = tl.load(a_ptr + offsets, mask=mask)
    b = tl.load(b_ptr + offsets, mask=mask)
    c = a + b
    tl.store(c_ptr + offsets, c, mask=mask)

如果你代码仍然是 mask = offsets < 15：

注：这段代码只适用于 N=15 的情况。

如果你改成 mask = offsets < N，并传入 N ，只要 \(N\le16\)，就可以用一个 program 正确处理。

那么这种解决方案的缺点在哪呢？

这种单 program 方案的缺点是， tl.arange(0, BLOCK_SIZE) 生成的 block tensor 大小不能无限大，Triton 对单个 block tensor 的元素数量有上限，例如常见上限是 1048576。

更重要的是，即使没有这个硬限制，让一个 program 处理过多元素也会导致寄存器压力过大、并行度不足、性能很差。因此不能靠无限增大 BLOCK_SIZE 来处理任意长度的向量。

这又应该怎么解决呢？

多 Block/Program 共同处理

与其受到单 Block 计算能力的限制，不如多 Block 共同分担计算，GPU 可以同时调度多个 CTA/thread block，从而获得更高的并行度。

在 Triton 中，grid 描述要启动多少个 program。它类似于 CUDA 中由多个 thread block 组成的 grid。

对于一维向量加法，只需要使用一维 grid，也就是只沿 \(x\) 方向排列 program 便已经足够了，在内核中我们可以使用 tl.program_id(axis=0) 来获取当前 program 在第 0 维 grid 上的编号。

若引入多 program，我们需要用 pid 来控制偏移。我们每个Block依旧只覆盖 16 （或 BLOCK_SIZE ）个逻辑位置，需要的 Block 数就是 ceil(N / 16) 或 ceil(N / BLOCK_SIZE) 。我们可以调用 triton 的 triton.cdiv(N, 16) 来计算。这样，如果想要获取指向正确 Block 的、正确位置的元素下标，就要这样算：

每个 program 需要先计算自己负责区间的起始下标

python 复制代码 
block_start = pid * BLOCK_SIZE

再加上

python 复制代码 
tl.arange(0, BLOCK_SIZE)

就能得到本 program 内所有元素的全局下标。

这里我们将 BLOCK_SIZE （之前都直接写完 \(16\)）参数化，得到这样的代码：

python 复制代码 
@triton.jit
def vector_add_kernel(a_ptr, b_ptr, c_ptr, N, BLOCK_SIZE: tl.constexpr[int]):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    a = tl.load(a_ptr + offsets, mask=mask)
    b = tl.load(b_ptr + offsets, mask=mask)
    c = a + b
    tl.store(c_ptr + offsets, c, mask=mask)
    
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
    BLOCK_SIZE = 16
    grid = (triton.cdiv(N, BLOCK_SIZE), )
    vector_add_kernel[grid](a, b, c, N, BLOCK_SIZE=BLOCK_SIZE)

至于 BLOCK_SIZE 应如何取值， BLOCK_SIZE 的最佳取值与数据规模、GPU 架构、num_warps 、寄存器压力、occupancy 和内存访问模式有关。实际工程中通常通过 benchmark 或 triton.autotune 来选择。

【Exercise】

查看 LeetGPU · 两数之和 测验你的本节的知识理解

热门推荐
01GitHub 镜像站点02Codex 接入 DeepSeek API 完整配置文档03CC-Switch & Claude 基于 Linux 服务器安装使用指南04零基础教你claude code 接入 deepseek V405【AI】2026 年具身智能模型和世界模型总结06Windows端Codex接入第三方模型（DeekSeek，BaiLian)07Cursor 接入 DeepSeek‑V4‑Pro 完整指南（2026 实测）08要裂开了！ChatGPT要手机号验证了？注册Codex要求验证电话号码怎么办？2026年登陆Codex要手机号验证的解决办法09codex app每次打开重连5次Reconnecting问题解决10裂开！ChatGPT 居然开始要手机号验证，附详细解决方法