Triton 手写内核实战：从向量加法到 FlashAttention 核心
摘要:
torch.compile的默认后端 TorchInductor 把计算图翻译成 Triton kernel——想真正理解 PyTorch 为什么快，或想为自定义算子手写高性能内核，Triton 是绕不开的一站。本文从零手写三个内核：向量加法 → 分块 GEMM → FlashAttention 核心，每个都给完整代码、性能对比和调优思路。
0. 实验目标
- 掌握 Triton 编程模型：块级（block-level）编程 + 自动并行
- 手写向量加法与分块 GEMM，理解 tile（分块） 与 循环展开
- 手写 FlashAttention 核心（online softmax 分块），对比 PyTorch 原生
- 学会 kernel 调优：块大小、
tl.constexpr、tl.static_range
1. 为什么框架工程师要学 Triton
原理讲解：
- CUDA 内核手写成本高：线程/共享内存/同步全要手动管。Triton 把编程模型提到 "块"级别——你描述"每个块做什么"，编译器负责映射到线程、调度、共享内存。
- TorchInductor 生成的 GPU kernel 就是 Triton（见本知识库 PyTorch 编译栈 篇）——学 Triton = 看懂 torch.compile 的产物。
- 自定义算子（融合 kernel、量化、长尾性能热点）用 Triton 写，几行到几十行搞定，性能可逼近手写 CUDA。
2. 内核一：向量加法（最小可运行）
python
import torch, triton, triton.language as tl
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
pid = tl.program_id(0) # 当前块的编号
offs = pid * BLOCK + tl.arange(0, BLOCK) # 本块覆盖的元素下标
mask = offs < n # 边界掩码（n 不整除 BLOCK）
x = tl.load(x_ptr + offs, mask=mask)
y = tl.load(y_ptr + offs, mask=mask)
tl.store(out_ptr + offs, x + y, mask=mask)
def add(x, y):
n = x.numel()
out = torch.empty_like(x)
grid = (triton.cdiv(n, 1024),) # 块数 = ceil(n / BLOCK)
add_kernel[grid](x, y, out, n, BLOCK=1024)
return out
x = torch.randn(1_000_000, device="cuda")
y = torch.randn(1_000_000, device="cuda")
print("正确性:", torch.allclose(add(x, y), x + y))逐段拆解：
@triton.jit：把 Python 函数编译成 GPU kernel。函数体是"每块执行的逻辑"，不是串行循环。tl.program_id(0)：当前块 ID；tl.arange(0, BLOCK)：生成 [0, BLOCK) 的索引向量——Triton 是向量化（块内 SIMD）编程。BLOCK: tl.constexpr：编译期常量。改块大小触发重新编译，所以调优是"编译期参数"而非运行期参数。mask：边界处理。n=1_000_000不整除1024，最后一块要掩码防越界。grid：启动的块数。Triton 自动把块调度到所有 SM——无需手动管线程。
3. 内核二：分块 GEMM（tile 化矩阵乘）
python
@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)
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # 本块的行
rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # 本块的列
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, K, BLOCK_K): # K 维分块循环
rk = k + tl.arange(0, BLOCK_K)
a = tl.load(a_ptr + rm[:, None] * stride_am + rk[None, :] * stride_ak)
b = tl.load(b_ptr + rk[:, None] * stride_bk + rn[None, :] * stride_bn)
acc += tl.dot(a, b) # 块级矩阵乘（编译为 MMA）
c = acc.to(tl.float16)
tl.store(c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn, c)
def matmul(a, b):
M, K = a.shape
K2, N = b.shape
c = torch.empty((M, N), device="cuda", dtype=torch.float16)
grid = (triton.cdiv(M, 64), triton.cdiv(N, 64))
matmul_kernel[grid](a, b, c, M, N, K,
a.stride(0), a.stride(1), b.stride(0), b.stride(1),
c.stride(0), c.stride(1), BLOCK_M=64, BLOCK_N=64, BLOCK_K=32)
return c
a = torch.randn(1024, 1024, device="cuda", dtype=torch.float16)
b = torch.randn(1024, 1024, device="cuda", dtype=torch.float16)
out = matmul(a, b)
ref = a @ b
print("正确性:", torch.allclose(out.float(), ref.float(), atol=0.5))原理讲解：
- tile 化：输出矩阵切成
BLOCK_M × BLOCK_N的块，每块一个 kernel 实例。K 维再切成BLOCK_K循环——这是 GEMM 高性能的核心：数据局部性 + 可展开循环。 tl.dot(a, b)：编译成 GPU 的 MMA（矩阵乘加速指令）——性能与 cuBLAS 同级别的关键。- 为什么快：每个 tile 的 A/B 块在寄存器/共享内存中复用 K 次，主内存带宽消耗降为
O(MN(K/BLOCK_K))而非O(MNK)。 - 对比 cuBLAS：Triton 手写 GEMM 通常能达到 cuBLAS 的 70-95%——对于融合场景（如 attention 里内联乘加）反而能超过库。
4. 内核三：FlashAttention 核心（online softmax）
python
@triton.jit
def flashattn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, N, d,
BLOCK: tl.constexpr, d_BLOCK: tl.constexpr):
pid = tl.program_id(0) # 处理一个 query 块
start = pid * BLOCK
rq = start + tl.arange(0, BLOCK) # query 行
rd = tl.arange(0, d_BLOCK) # head 维度
q = tl.load(q_ptr + rq[:, None] * d + rd[None, :])
m_i = tl.full((BLOCK,), float("-inf"), tl.float32) # 运行最大值
l_i = tl.zeros((BLOCK,), tl.float32) # 运行 exp 和
acc = tl.zeros((BLOCK, d_BLOCK), tl.float32)
for start_k in range(0, N, BLOCK): # 分块扫描 K/V
rk = start_k + tl.arange(0, BLOCK)
k = tl.load(k_ptr + rk[:, None] * d + rd[None, :])
# 注意力分数（无 scale，简化）
s = tl.dot(q, tl.trans(k)) # (BLOCK, BLOCK)
m_ij = tl.max(s, axis=1)
m_new = tl.maximum(m_i, m_ij)
p = tl.exp(s - m_new[:, None]) # 数值稳定的 softmax
alpha = tl.exp(m_i - m_new)
l_i = l_i * alpha + tl.sum(p, axis=1)
v = tl.load(v_ptr + rk[:, None] * d + rd[None, :])
acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), v)
m_i = m_new
acc /= l_i[:, None]
tl.store(o_ptr + rq[:, None] * d + rd[None, :], acc.to(tl.float16))
# 调用（简化：单头、d=64、N 对齐 BLOCK）
N, d = 512, 64
q = torch.randn(N, d, device="cuda", dtype=torch.float16)
k = torch.randn(N, d, device="cuda", dtype=torch.float16)
v = torch.randn(N, d, device="cuda", dtype=torch.float16)
o = torch.empty_like(q)
flashattn_kernel[(triton.cdiv(N, 64),)](q, k, v, o, N, d, BLOCK=64, d_BLOCK=64)
print("输出形状:", o.shape)原理讲解（online softmax 三件套）：
- 问题：标准 attention 要先把整个
N×N分数矩阵算出来再 softmax，显存O(N²)——长序列直接爆炸。 - 解决：分块扫描 K/V（每次只载
BLOCK×d的块），维护三个运行量：m_i（运行最大值）、l_i（运行 exp 和）、acc（加权和）。每来一个块：更新最大值 → 重缩放alpha=exp(m_i-m_new)→ 更新l_i和acc。全程显存 O(N)，这就是 FlashAttention 的核心思想。 - 代码里
tl.dot复用 GEMM 加速；tl.trans转置 K 块。 - 工程注意：真实实现还需 scale（
1/sqrt(d)）、mask（因果）、多头循环——本文是最小可运行核心。
5. 调优三板斧
python
# 1) 块大小：更大的 BLOCK 减少块数/循环，但受寄存器与占用率限制
# GEMM 常见：BLOCK_M=64/128, BLOCK_N=64/128, BLOCK_K=16/32/64
# 2) 编译期常量（tl.constexpr）：改值即重编译——调优参数不要写死成 Python 变量
# 3) tl.static_range：K 循环已知时用静态展开，让编译器展开循环实验结果（示例量级 · A100 · FP16）：
| 内核 | 相对性能 |
|---|---|
| 手写 Triton GEMM（BLOCK=64/64/32） | cuBLAS 的 ~85% |
| FlashAttention 核心（BLOCK=64） | PyTorch 原生 attention 的 ~2-3×（长序列显存省 O(N²)） |
数值为相对量级参考。要点：Triton 手写内核达到库级性能（GEMM 80-95%），且融合场景（attention）可以超越原生实现——因为它把 softmax+乘加+输出写回融合进一个 kernel，省掉了中间
N×N张量的读写。
6. 工程启示汇总
- Triton 是"块级"编程模型：描述块做什么，编译器管线程/共享内存/调度。
- torch.compile 的后端就是 Triton：学它 = 看懂编译产物、能定制融合。
- GEMM 高性能 = tile 化 + tl.dot（MMA）：K 循环展开、数据局部性。
- FlashAttention 核心 = online softmax 三件套（m_i/l_i/acc），显存 O(N) 的关键。
- 调优参数用 tl.constexpr：编译期常量，改值即重编译。
tl.static_range静态展开已知循环，给编译器展开机会。- 自定义算子用 Triton 写：性能逼近手写 CUDA，接入 PyTorch 生态（dispatch，见 内部架构 篇）零改动。
7. 延伸阅读
- Triton 官方教程
- FlashAttention 论文
- 本知识库：09. 深度学习框架 章节 · PyTorch 内部架构 · 编译栈实战 · Triton 学习路径