目录

  1. Triton 是什么
  2. 核心心智模型:块级 SPMD
  3. 和 CUDA 的对照
  4. 第一个内核:向量加法(逐行讲解)
  5. tl 语言核心 API 速查
  6. program_id / arange / 偏移 / mask 四件套
  7. constexpr:编译期常量到底是什么
  8. 指针、stride 与多维张量寻址
  9. load / store 与 mask 边界处理
  10. 归约、点积与跨轴运算
  11. num_warps / num_stages / BLOCK 与 autotune
  12. 进阶例子:融合 softmax
  13. 进阶例子:矩阵乘法骨架
  14. SPMD 之外:原子操作与指针表
  15. 调试技巧
  16. 常见坑
  17. 一页纸总结

1. Triton 是什么

Triton 是一个用 Python 写 GPU 内核(kernel)的语言 + 编译器,由 OpenAI 开源。你用 @triton.jit 装饰一个 Python 函数,Triton 把它编译成 GPU 机器码(NVIDIA 上经 LLVM → PTX → SASS)。

为什么用它而不是直接写 CUDA C++:

维度CUDA C++Triton
语言C++Python 子集
编程粒度单线程(thread)视角单”块”(block)视角,块内自动并行
内存合并/bank 冲突手动管理编译器大量自动处理
共享内存手动 __shared__编译器自动调度(大多数情况)
自动调优手写多版本@triton.autotune 内建
上手难度高中
能压榨到的极限性能最高接近(很多场景持平 cuBLAS/cuDNN)

一句话:Triton 让你用接近 NumPy 的写法,拿到接近手写 CUDA 的性能,代价是放弃一部分极端底层的控制。


2. 核心心智模型:块级 SPMD

这是理解 Triton 的唯一最重要的点。

  • 你启动一个 grid(网格),grid 由很多个 program instance(程序实例) 组成。
  • 每个 program instance 跑的是同一份内核代码(SPMD = Single Program Multiple Data),但通过 tl.program_id() 知道”我是第几个”,从而处理不同的数据切片。
  • 关键差异:在 CUDA 里你写”一个线程处理一个元素”;在 Triton 里你写”一个 program 处理一整块(BLOCK)元素”,块内的线程/warp 由编译器自动安排,你看不到也不用管。

图示(假设处理一个长度 1024 的数组,BLOCK=256):

grid = (4,)   # 启动 4 个 program
 
program_id=0 → 处理元素 [0   : 256)
program_id=1 → 处理元素 [256 : 512)
program_id=2 → 处理元素 [512 : 768)
program_id=3 → 处理元素 [768 :1024)
 
每个 program 内部:对它那 256 个元素做"向量化"运算(编译器把这 256 路摊到 warp 上)

所以你写内核时的思维是:“假设我是其中一个 program,我负责哪一块?这一块怎么算?“


3. 和 CUDA 的对照

帮有 CUDA 背景的人快速迁移:

概念CUDATriton
块索引blockIdx.xtl.program_id(0)
块数量gridDim.xtl.num_programs(0)(或启动时的 grid)
线程索引threadIdx.x没有直接对应(块内自动)
块内偏移向量手动 threadIdx.x 循环tl.arange(0, BLOCK)
读全局内存x[i]tl.load(x_ptr + offs, mask=...)
写全局内存x[i] = vtl.store(x_ptr + offs, v, mask=...)
共享内存__shared__编译器自动(一般不用手写)
同步__syncthreads()块内自动(一般不用手写)
编译期常量模板参数 template<int N>N: tl.constexpr

没有 threadIdx 是 Triton 最大的”不适应点”:你不再操心单个线程,而是操心一整块向量。


4. 第一个内核:向量加法(逐行讲解)

经典的 “hello world”:z = x + y。

import torch
import triton
import triton.language as tl
 
@triton.jit
def add_kernel(
    x_ptr,            # 输入 x 的首地址(指针)
    y_ptr,            # 输入 y 的首地址
    z_ptr,            # 输出 z 的首地址
    n_elements,       # 元素总数(运行期变量)
    BLOCK_SIZE: tl.constexpr,   # 每个 program 处理多少元素(编译期常量)
):
    # 1) 我是第几个 program?
    pid = tl.program_id(axis=0)
 
    # 2) 我负责的这一块,起始下标 = pid * BLOCK_SIZE
    block_start = pid * BLOCK_SIZE
 
    # 3) 这一块覆盖的所有下标(一个向量,长度 BLOCK_SIZE)
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
 
    # 4) 越界保护:最后一块可能不满 BLOCK_SIZE
    mask = offsets < n_elements
 
    # 5) 从全局内存读两块数据(被 mask 掉的位置不读,避免越界)
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
 
    # 6) 向量化计算(这一句作用在整块 BLOCK_SIZE 个元素上)
    z = x + y
 
    # 7) 写回(同样用 mask 保护)
    tl.store(z_ptr + offsets, z, mask=mask)
 
 
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    z = torch.empty_like(x)
    n = x.numel()
    # grid:需要多少个 program。triton.cdiv(a,b) = 向上取整除法
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
    add_kernel[grid](x, y, z, n, BLOCK_SIZE=1024)
    return z

几个要点:

  • add_kernel[grid](...):方括号里是 grid(启动多少个 program),圆括号里是内核参数。
  • grid 写成 lambda 是因为它依赖 BLOCK_SIZE(可被 autotune 改写),用 meta 拿到最终值。
  • Torch tensor 传进去会自动退化成它的 data_ptr()。
  • tl.arange(0, BLOCK_SIZE) 要求 BLOCK_SIZE 编译期已知,所以它是 constexpr。

5. tl(triton.language)核心 API 速查

import triton.language as tl,最常用的:

API作用
tl.program_id(axis)当前 program 在 grid 第 axis 维的索引(0/1/2)
tl.num_programs(axis)grid 第 axis 维的总数
tl.arange(start, end)编译期定长向量 [start, ..., end-1]
tl.load(ptr, mask, other)从全局内存读;mask=False 的位置取 other(默认 0)
tl.store(ptr, val, mask)写回全局内存;mask=False 的位置不写
tl.zeros(shape, dtype)造一个全 0 的块(shape 必须编译期已知)
tl.full(shape, value, dtype)造一个填充值的块
tl.where(cond, a, b)逐元素三元选择
tl.maximum / tl.minimum逐元素 max/min
tl.exp / tl.log / tl.sqrt / tl.sigmoid逐元素数学函数
tl.sum / tl.max / tl.min(x, axis)沿某轴归约
tl.dot(a, b)矩阵乘(用 Tensor Core)
tl.cdiv(a, b)向上取整除(也有 triton.cdiv)
tl.static_range(n)编译期展开的 for 循环
tl.atomic_add(ptr, val, mask)原子加(多 program 写同一地址时用)
tl.constexpr类型标注:声明参数为编译期常量
x.to(tl.float32)dtype 转换

6. program_id / arange / 偏移 / mask 四件套

这四个几乎出现在每个内核里,组合套路固定。以二维为例(处理一个 [M, N] 矩阵,每个 program 负责 BLOCK_M × BLOCK_N 的小块):

@triton.jit
def tile_kernel(ptr, M, N,
                stride_m, stride_n,
                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    # 1) 我在二维 grid 里的位置
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
 
    # 2) 本块覆盖的行号、列号(各是一个向量)
    rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)   # 形状 (BLOCK_M,)
    cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)   # 形状 (BLOCK_N,)
 
    # 3) 越界 mask(行、列各自一份,再广播成二维)
    row_mask = rows < M           # (BLOCK_M,)
    col_mask = cols < N           # (BLOCK_N,)
    mask = row_mask[:, None] & col_mask[None, :]      # (BLOCK_M, BLOCK_N)
 
    # 4) 二维偏移:用 [:, None] / [None, :] 做广播,生成 (BLOCK_M, BLOCK_N) 偏移矩阵
    offs = rows[:, None] * stride_m + cols[None, :] * stride_n
 
    # 5) 读 / 算 / 写
    block = tl.load(ptr + offs, mask=mask, other=0.0)
    block = block * 2.0
    tl.store(ptr + offs, block, mask=mask)

记牢这套范式:

索引向量 = program_id * BLOCK + tl.arange(0, BLOCK)
mask     = 索引向量 < 真实尺寸
偏移     = 行索引[:,None]*行stride + 列索引[None,:]*列stride
读       = tl.load(base_ptr + 偏移, mask=mask, other=...)
写       = tl.store(base_ptr + 偏移, 值, mask=mask)

[:, None] 把 (BLOCK_M,) 变成 (BLOCK_M, 1),[None, :] 把 (BLOCK_N,) 变成 (1, BLOCK_N),两者相加触发广播,得到 (BLOCK_M, BLOCK_N)。这和 NumPy 完全一样。


7. constexpr:编译期常量到底是什么

tl.constexpr 是初学者最容易误解的点,单独讲清楚。

它是”修饰符”,不是”值”

X: tl.constexpr 的意思是:参数 X 的值在编译内核时就必须确定(即编译期常量),而不是运行时才传入的普通数据。它类似 C++ 模板参数 template<int X>。

所以当你看到:

def kernel(src, NUM_ROWS: tl.constexpr, NUM_COLS: tl.constexpr, ...):

NUM_ROWS 和 NUM_COLS 不是”同一个东西”,更不是”值相同”。它们是两个不同的量(行数 vs 列数),只是都恰好在编译期已知而已。constexpr 只标注”何时确定”,不标注”是什么值”。

constexpr vs 普通运行期参数

@triton.jit
def k(x_ptr,                    # 普通参数:指针,运行期传
      n_rows,                   # 普通参数:int 标量,运行期传(改值不重新编译)
      BLOCK: tl.constexpr):     # 编译期常量:改值会触发重新编译
    ...
类型何时确定改变它的代价能否用于形状/static_range
普通参数(指针/标量)运行期几乎免费,不重编译否
tl.constexpr编译期每个新值都重新编译一份内核可以

为什么要把某些参数设成 constexpr

  1. 形状必须编译期已知:tl.arange(0, BLOCK)、tl.zeros((BLOCK_M, BLOCK_N))、tl.static_range(K) 都要求参数是 constexpr,因为编译器要据此分配寄存器/展开循环。
  2. 特化(specialization)优化:值固定后,编译器可以
    • 把循环完全展开(tl.static_range)
    • 在某些场景消除 mask(比如 N 恰好是 BLOCK 的整数倍时不需要边界判断)
    • 把除法/取模优化成移位/与运算
    • 更好地分配寄存器和共享内存
  3. stride 常量化:把 stride 设成 constexpr,编译器能把寻址算式里的乘法折叠掉。

代价

每出现一组新的 constexpr 取值组合,Triton 就会重新编译一份内核(并缓存)。所以 constexpr 适合”取值种类不多”的量(BLOCK 大小、维度档位),不适合”几乎每次都不同”的量(否则编译缓存爆炸)。


8. 指针、stride 与多维张量寻址

Triton 里没有”多维数组下标 a[i][j]”这种语法,一切都是一维指针 + 偏移。多维通过 stride(步长) 表达,和 PyTorch 的 stride 概念一致。

对一个 [M, N] 行主序(row-major)矩阵:

元素 (i, j) 的线性偏移 = i * stride_m + j * stride_n
行主序时:stride_m = N, stride_n = 1

向量化版本(一次算一个块):

offs = rows[:, None] * stride_m + cols[None, :] * stride_n   # (BLOCK_M, BLOCK_N) 偏移矩阵
val  = tl.load(base_ptr + offs, mask=mask)

为什么传 stride 而不是直接传形状:这样内核能直接支持转置、切片、非连续张量(只要把对应的 stride 传进来即可),不必要求输入连续。PyTorch 里 t.stride(0)、t.stride(1) 就能拿到。


9. load / store 与 mask 边界处理

x = tl.load(ptr + offs, mask=mask, other=0.0)
tl.store(ptr + offs, val, mask=mask)
  • mask:布尔块,形状和 offs 一致。mask 为 True 的位置才真正访存,False 的位置:
    • load:不读内存,用 other(默认 0)填充;
    • store:不写内存(保护越界)。
  • 为什么必须有 mask:数据长度通常不是 BLOCK 的整数倍,最后一个 program 的块会”探出边界”。没有 mask 就会越界读写 → 非法内存访问 / 数据损坏。
  • other:被 mask 掉的 load 返回值。做 max 归约时常设成 -inf,做 sum 时设 0。

性能提示:当某维度恰好是 BLOCK 的整数倍、且尺寸是 constexpr 时,编译器能把 mask 优化掉,所以”对齐的尺寸”往往更快。


10. 归约、点积与跨轴运算

# 沿 axis 归约(把那一维压掉)
row_sum = tl.sum(x, axis=1)        # x:(BLOCK_M, BLOCK_N) → (BLOCK_M,)
row_max = tl.max(x, axis=1)        # 同上
 
# 矩阵乘:走 Tensor Core,a:(BM,BK) b:(BK,BN) → (BM,BN)
acc = tl.dot(a, b)                 # 通常累加到 fp32 acc
 
# 编译期展开的循环(K 维分块累加常用)
acc = tl.zeros((BM, BN), dtype=tl.float32)
for k in tl.static_range(0, K, BK):
    a = tl.load(...); b = tl.load(...)
    acc += tl.dot(a, b)
  • tl.dot 要求输入是 2D 块,且形状满足 Tensor Core 约束(一般维度 ≥16)。
  • 累加器(accumulator)通常用 fp32,即使输入是 fp16/bf16,以保证数值精度。
  • 归约的 axis 是块内的轴,不是 grid 的轴。

11. num_warps / num_stages / BLOCK 与 autotune

这些是性能旋钮:

旋钮含义影响
BLOCK_*每个 program 处理的块大小太小→并行多但每个程序工作少、访存效率低;太大→寄存器/共享内存压力大,occupancy 下降
num_warps每个 program 用几个 warp(每 warp 32 线程)决定块内并行宽度;常见 4/8/16
num_stages软件流水线级数(load 与 compute 重叠)越大越能掩盖访存延迟,但占更多共享内存

手动指定:

add_kernel[grid](x, y, z, n, BLOCK_SIZE=1024, num_warps=4, num_stages=2)

自动调优(让 Triton 跑一遍挑最快配置并缓存):

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64},  num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64},  num_warps=8, num_stages=4),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=3),
    ],
    key=["M", "N"],   # 当 M/N 变化时,重新选择最优配置
)
@triton.jit
def kernel(...):
    ...

key 指定”哪些参数变化时需要重新 autotune”;同一组 key 的结果会被缓存。


12. 进阶例子:融合 softmax

按行做 softmax([M, N],每个 program 处理一行),展示归约 + 数值稳定:

@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_row_stride, out_row_stride,
                   n_cols, BLOCK_N: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_N)
    mask = cols < n_cols
 
    # 读一整行(越界处填 -inf,保证不影响 max)
    x = tl.load(in_ptr + row * in_row_stride + cols, mask=mask, other=-float("inf"))
 
    # 数值稳定:先减去行最大值,再 exp
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    den = tl.sum(num, axis=0)
    y = num / den
 
    tl.store(out_ptr + row * out_row_stride + cols, y, mask=mask)

启动:

grid = (M,)   # 每行一个 program
softmax_kernel[grid](out, x, x.stride(0), out.stride(0), N,
                     BLOCK_N=triton.next_power_of_2(N))

要点:BLOCK_N 取 ≥N 的 2 的幂,这样一行能一次性装进一个块;越界用 -inf 填充保证 max/exp 正确。这就是经典的”融合”——读、减最大值、exp、求和、归一化、写回,全部在一个内核、一次内存往返里完成,省掉中间张量。


13. 进阶例子:矩阵乘法骨架

C[M,N] = A[M,K] @ B[K,N],展示 K 维分块累加 + tl.dot:

@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 块、B 块的起始指针(二维偏移)
    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
 
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k, other=0.0)
        acc += tl.dot(a, b)               # Tensor Core 累加
        a_ptrs += BLOCK_K * stride_ak     # 指针沿 K 推进
        b_ptrs += BLOCK_K * stride_bk
 
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=c_mask)

核心思想:把大矩阵切成 BLOCK_M × BLOCK_N 的输出瓦片(tile),每个 program 算一个瓦片;K 维再切成 BLOCK_K 一段段地 tl.dot 累加。指针随 K 推进(a_ptrs += BLOCK_K * stride_ak)是性能关键写法。


14. SPMD 之外:原子操作与指针表

原子操作

多个 program 可能写同一个地址(比如做直方图、scatter-add),这时需要原子操作避免竞争:

tl.atomic_add(out_ptr + idx, val, mask=mask)   # 原子加
tl.atomic_max(out_ptr + idx, val, mask=mask)   # 原子取最大

指针表(pointer table)

Triton 支持把”一组指针”作为输入(一个 int64 张量,每个元素是一个地址),内核内部 tl.load 出某个地址再当指针用。这让一个内核能写到多个不同的目标缓冲区:

@triton.jit
def scatter_to_buffers(src, dst_ptrs,           # dst_ptrs: 指针表(int64)
                       dst_ids,                  # 每行去哪个 buffer
                       NUM_COLS: tl.constexpr, BLOCK_N: tl.constexpr):
    row = tl.program_id(0)
    col = tl.arange(0, BLOCK_N)
    mask = col < NUM_COLS
 
    values   = tl.load(src + row * NUM_COLS + col, mask=mask)
    dst_base = tl.load(dst_ptrs + tl.load(dst_ids + row))   # 取出目标 buffer 的基址
    dst_ptr  = dst_base.to(tl.pointer_type(tl.float32)) + col
    tl.store(dst_ptr, values, mask=mask)

关于”写到远端内存”的原理(重要概念):tl.store 本身只是”往某个地址写”。这个地址指向哪里、物理上在不在本设备,取决于这个地址是怎么来的:

  • 如果地址是本 GPU 的显存 → 普通写。
  • 如果地址是对端 GPU 显存被映射进本进程地址空间(例如通过 CUDA IPC / NVLink P2P,或 PyTorch SymmetricMemory 的 CUDA 后端)→ 这次 tl.store 会被互连硬件(NVLink/NVSwitch)直接 DMA 到对端 GPU,不经过网卡、不跨机。这是 load/store 语义的远程访问。
  • 如果要**跨节点(跨机)**写到别的机器的 GPU → 普通 tl.store 做不到,需要 NVSHMEM 这类对称内存运行时:地址是 NVSHMEM 对称地址,device 端运行时发现目标在远端节点时,会触发 GPU-initiated RDMA(IBGDA),由 GPU 直接驱动 InfiniBand 网卡完成 RDMA write。这条路才会落到 IB verbs(ibv_post_send 一类)。

一句话:tl.store 永远只是”写一个地址”;能不能跨卡/跨机、走 NVLink 还是走 IB,完全由那个地址背后的内存映射机制决定,而不是 store 指令本身。


15. 调试技巧

手段说明
TRITON_INTERPRET=1设此环境变量,内核在 CPU 解释器上单步跑,可以 print、用 pdb,定位逻辑错误(慢,但能调试)
tl.device_print("msg", x)在内核里打印张量值(GPU 上)
tl.static_assert(cond)编译期断言
先用小 shape用 M=N=8 之类的小尺寸,和 PyTorch 参考实现对拍
torch.testing.assert_close把 Triton 结果和 torch 等价实现比对
看 PTX/SASSkernel.asm["ptx"] 可拿到生成的 PTX,排查性能问题

对拍模板:

out_triton = my_triton_op(x)
out_ref = torch_reference(x)
torch.testing.assert_close(out_triton, out_ref, rtol=1e-3, atol=1e-3)

16. 常见坑

  1. 忘记 mask → 越界访问,结果随机错或 CUDA illegal memory access。任何”尺寸不一定整除 BLOCK”的访存都要 mask。
  2. constexpr 用成运行期变量 → 报错或性能差。形状相关的量(BLOCK、用于 arange/zeros 的尺寸)必须是 constexpr。
  3. constexpr 取值太发散 → 每个新值重编译,首次运行巨慢、缓存膨胀。只把”档位有限”的量设 constexpr。
  4. accumulator 用低精度 → fp16 累加丢精度;tl.dot 的累加器用 fp32。
  5. stride 写死成连续 → 输入非连续(转置/切片)时结果错。老老实实传 tensor.stride(i)。
  6. grid 算错 → 用 triton.cdiv(n, BLOCK) 向上取整,别用整除漏掉尾块。
  7. 以为有 threadIdx → Triton 是块级编程,没有线程级索引;想”每线程一元素”的写法要改成”每 program 一块向量”。
  8. tl.arange 长度非 2 的幂 → 多数情况要求 2 的幂,用 triton.next_power_of_2()。
  9. 在内核里写依赖运行期张量值的 Python 控制流 → if x > 0(x 是块)不行,要用 tl.where;只有 constexpr 条件才能用真正的 if。

17. 一页纸总结

心智模型:块级 SPMD —— 你写"一个 program 处理一块数据",grid 决定有多少 program。
 
固定四件套:
  pid     = tl.program_id(axis)
  offsets = pid * BLOCK + tl.arange(0, BLOCK)
  mask    = offsets < size
  data    = tl.load(ptr + offsets, mask=mask, other=0)
  ...算...
  tl.store(ptr + offsets, result, mask=mask)
 
多维:用 stride 寻址 + [:,None]/[None,:] 广播
  offs = rows[:,None]*stride_m + cols[None,:]*stride_n
 
constexpr:编译期常量(≈模板参数),决定形状、触发特化优化;不是"值",只是"何时确定"。
  - 形状/arange/static_range/zeros 的尺寸必须 constexpr
  - 取值档位要少,否则重编译爆炸
 
性能旋钮:BLOCK 大小 / num_warps / num_stages,用 @triton.autotune 自动挑。
 
归约用 tl.sum/tl.max(x, axis=...);矩阵乘用 tl.dot(累加器 fp32)。
 
store 只是"写地址";跨卡靠 NVLink P2P 映射,跨机靠 NVSHMEM/IBGDA RDMA —— 取决于地址的内存映射,不取决于 store 本身。
 
调试:TRITON_INTERPRET=1 上 CPU 跑 + 和 torch 对拍。

推荐学习路径

  1. 跑通本文第 4 节的向量加法。
  2. 改写第 12 节的 softmax,和 torch.softmax 对拍。
  3. 啃第 13 节的 matmul,理解 K 维分块和指针推进。
  4. 给 matmul 加 @triton.autotune,对比不同 BLOCK 的性能。
  5. 读官方 tutorials(triton-lang.org 的 python/tutorials/),里面有 fused-attention、layer-norm、dropout 等完整例子。