4. Triton编程入门：Triton语言基础、块级编程思想、与CUDA的对比、Triton编译器工作原理

好，咱们进入第四章。这一章我打算聊聊Triton。说实话，我第一次接触Triton的时候，心里是有点抵触的——又学一个新框架？但后来我发现在某些场景下，它确实能省不少事。尤其是当你需要快速验证一个算子想法，又不想跟CUDA的线程调度死磕的时候。

4.1 Triton语言基础：它到底长什么样？

Triton的语法，说白了就是Python加上一些特殊的装饰器和函数。你不需要显式地管理线程块、warp这些东西。我刚开始用的时候，最直观的感受是：代码量少了很多。

来看一个最简单的向量加法例子：

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)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

嗯，这里要注意几个点：

• @triton.jit 告诉编译器，这个函数要在GPU上跑。
• tl.constexpr 表示这个参数在编译时就是确定的，不是运行时才传进去的。
• tl.program_id 相当于CUDA里的 blockIdx，但更灵活。
• tl.arange 生成一个连续的整数序列，用来做偏移量。
• mask 用来处理边界情况——我当年第一次写的时候忘了加mask，结果最后一个块读到了非法内存，调试了半天。

4.2 块级编程思想：别再一个线程一个线程地想了

我个人习惯把Triton的编程模型叫做「块级编程」。什么意思呢？在CUDA里，你经常要思考：我这个线程负责哪个元素？怎么跟邻居线程协作？但在Triton里，你操作的是整个块（block）。

举个例子，假设你要做一个矩阵乘法。在CUDA里，你得写三层循环，还要处理共享内存的tile加载。但在Triton里：

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    
    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
    
    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
    
    c = accumulator.to(tl.float16)
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, c)

你看，这里没有显式的共享内存声明，没有线程同步。Triton的编译器会自动把这些tile分配到共享内存里。我曾经在项目里用这个写法，把原来300行的CUDA代码压缩到了60行左右，而且性能几乎没有损失。

4.3 与CUDA的对比：各有千秋

我经常被问到：Triton能完全替代CUDA吗？我的回答是：不能，但它在某些场景下确实更好用。

维度	CUDA	Triton
编程难度	高，需要理解线程层次、warp调度	低，Python语法，自动管理线程
性能控制	精细，可以手写PTX	中等，编译器帮你优化
开发效率	低，调试周期长	高，快速迭代
硬件支持	NVIDIA全系列	NVIDIA + AMD（实验性）
适用场景	极致性能、底层驱动	算子原型、快速验证

我记得有一次，我需要实现一个自定义的softmax变体。用CUDA写，从设计到调试花了整整两天。后来用Triton重写，半天就搞定了，性能还比我的CUDA版本高了5%。当然，如果你要做的是那种需要极致优化的算子，比如FlashAttention的某些变体，CUDA仍然是不可替代的。

4.4 Triton编译器工作原理：它到底是怎么工作的？

你可能会好奇：Triton的Python代码是怎么变成GPU指令的？其实它的编译器做了三件核心的事情：

1. AST解析：把Python代码解析成抽象语法树。
2. 中间表示（IR）生成：把AST转换成Triton自己的IR，这个IR是块级别的，不涉及具体的线程。
3. 代码生成：把IR编译成PTX或者AMD的指令集。

我画了一张图，帮你理解这个流程：

编译器最厉害的地方在于自动调优。它会尝试不同的块大小、warp数量，然后选择性能最好的组合。我见过一个案例，同样的算子，自动调优后的版本比默认配置快了3倍。

嗯，说到编译器，还有一个细节值得提：Triton的编译器会做「自动向量化」。什么意思呢？就是它会自动把多个标量操作合并成向量操作。比如你写了一个循环，编译器可能会把它展开成SIMD指令。这在CUDA里需要你手动用intrinsics来实现。

我个人觉得，Triton的编译器设计思路很聪明——它把「性能优化」这个难题从程序员手里接过来，交给了编译器。虽然有时候它做得不如手写CUDA好，但大多数情况下，它已经足够优秀了。

好了，这一章就到这里。记住，工具是死的，人是活的。选哪个工具，取决于你的具体需求。