Triton：块级编程

手写 CUDA 要管线程、占用率、warp——每个都是坑。Triton 把"线程层"抹掉，让你直接写"块"（tile）：声明一块数据怎么算，编译器负责把它映射到线程、warp、共享内存。用 Python 写，性能接近手写 CUDA——这是 FlashAttention 官方实现、PyTorch 内置算子、Unsloth 等大量项目的选择。

块级编程：少一层抽象

对比同一个向量加法：

// CUDA：显式管线程
__global__ void add(float* C, float* A, float* B, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) C[i] = A[i] + B[i];
}

# Triton：显式管块，编译器管线程
@triton.jit
def add_kernel(C, A, B, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)      # 一块索引
    mask = offs < n
    a = tl.load(A + offs, mask=mask)              # 一块加载
    b = tl.load(B + offs, mask=mask)
    tl.store(C + offs, a + b, mask=mask)          # 一块写回

区别一目了然：没有 thread、没有 blockDim、没有手动合并访存——tl.arange 声明一个块，编译器决定用多少线程、怎么对齐、怎么喂 Tensor Core。你只需要管三件事：

块多大（BLOCK 是编译期常量，自动 specialization）
块在网格的哪里（program_id）
块内怎么算（tl.load / 计算 / tl.store）

为什么它快

Triton 不是"慢的 Python 包装"，它是一个真正的编译器：

自动映射：块 → warp → 线程的映射、寄存器分配、占用率由编译器决定，且比手写更不容易犯错
自动访存优化：合并访存、向量化加载自动生成
块级融合：多个 tl.load + 计算天然融合成一个内核——算子融合是默认行为
Tensor Core 自动喂：tl.dot 自动映射到 wgmma，形状对齐、数据搬运由编译器安排

这就是"接近手写 CUDA"的来源：你写块级逻辑（决定算术强度），编译器写线程级逻辑（决定怎么打满硬件）——分工恰好落在各自最擅长的层。

一个真实的例子：softmax

Triton 官方的 fused softmax 例子，一个内核完成"读入 → max → exp → 归约 → 写出"（朴素 PyTorch 要 5 个内核 + 中间显存往返）：

@triton.jit
def softmax_kernel(X, Y, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    x = tl.load(X + row * n_cols + cols, mask=cols < n_cols, other=-float('inf'))
    x = x - tl.max(x, axis=0)                     # 数值稳定的 max
    num = tl.exp(x)
    den = tl.sum(num, axis=0)
    tl.store(Y + row * n_cols + cols, num / den, mask=cols < n_cols)

五个算子 → 一个内核，显存往返 ×1（原来是 ×5），对 4096 宽的行实测提速 2–4 倍——这就是算术强度一章"融合翻倍"的直接应用。

与 CUDA 的分工

	Triton	手写 CUDA
写法	Python，块级	C++，线程级
性能上限	接近手写（90%+）	极限
开发速度	快（小时级）	慢（天级）
控制粒度	块级（线程交给编译器）	完全控制
异步流水/TMA	部分支持，仍在进化	完全控制
调试	print 调试，工具较简单	成熟工具链

选型经验：

新算子原型/自适应逻辑（变长、动态 shape、特殊融合）：Triton 首选
标准算子：用库（cuBLAS/CUTLASS），别重造
极限性能、特殊硬件特性（TMA 流水、FP8 特殊路径）：手写 CUDA/CUTLASS
多硬件后端：Triton 的多后端（AMD、昇腾等）是它越来越重要的理由——一套 Triton 代码，多个后端跑

工程实践

BLOCK 大小是最重要的超参：从 128/256 起步，用 triton.testing 扫一遍；constexpr 参数会触发重新编译（动态 shape 注意缓存命中率）
num_warps / num_stages：每个 kernel 的并发与流水深度，默认 4/2 起步，访存密集型加大 num_stages
数值验证：与 PyTorch 参考实现对拍（容差按精度设置），mask 边界与 -inf 填充是最常见的 bug 源
性能验证：先 Roofline 算上限（融合后的算术强度），再对比实测——Triton 写出来慢，通常是块大小或 stages 没调，而不是框架慢

本章要点

Triton 抹掉线程层：写块级逻辑，编译器管线程/warp/共享内存/Tensor Core。
块级融合是默认行为；fused softmax 一例就是融合收益的直接展示。
性能接近手写（90%+），开发速度快一个数量级；极限控制仍需手写 CUDA。
BLOCK / num_warps / num_stages 是三个核心旋钮；先 Roofline 算上限再调。
多后端是战略价值：一套代码多硬件跑，国产芯片生态的常见入口。