让编码智能体写 Triton 算子：对拍验证 + 性能基准的完整闭环

这篇能做出什么

跑完这篇教程，你会得到一套可以反复使用的流水线：

1. 一句算子需求（本文用 AI 词典：RMSNorm">RMSNorm 前向做例子），交给编码智能体，让它产出 Triton 内核；

2. 一个随机对拍脚本，用 PyTorch 参考实现当“标准答案”，在多种 shape、dtype、数值量级上比对，输出最大绝对误差与最大相对误差；

3. 一个性能基准脚本，测延迟、算有效带宽、和 PyTorch 基线比加速比；

4. 一张判定表：达标就合并，不达标就回炉调参，仍然不行就放弃这条路线。

关键点不在于“让智能体写出一个能跑的 kernel”——那太容易了。关键在于先写测试、再写实现，让智能体的产出必须过对拍这一关；然后用带宽而不是用感觉来判断它值不值得留在代码库里。没有闭环，你得到的只是一段看起来很像 CUDA 的代码。

前置条件清单

• 一台有 GPU 的机器，能跑 CUDA。Triton 也支持部分其他加速器，具体支持矩阵以官方文档当前版本为准。
• Python 环境，装好 torch、triton。版本以官方文档当前版本为准，不要抄别人博客里的固定版本号。
• 一个能读写文件、能执行 shell 命令的编码智能体（IDE 插件或命令行形态都行）。
• 一个干净的 git 仓库。整个流程会产生多轮修改，没有版本控制你会分不清哪版是对的。
• 会看 nvidia-smi。

先确认环境可用：

```bash

python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

python -c "import triton; print(triton.__version__)"

```

如果 cuda.is_available() 是 False，后面所有步骤都跑不通，先解决驱动和 CUDA 环境。

第 1 步：把需求写成“可验收的规格”

智能体最怕的不是难题，是模糊。你说“写个 rmsnorm 算子”，它会给你一堆假设：eps 放哪、累加用什么精度、输入是不是连续、支不支持多维。这些假设里只要有一条和你的预期不同，对拍就会失败，而你会以为是 kernel 写错了。

所以先写一份规格，直接放进 prompt：

```text

目标：实现 RMSNorm 前向，Triton 内核。

输入：

• x: 形状 (M, N) 的连续张量，dtype 为 float32 / float16 / bfloat16
• w: 形状 (N,) 的权重，dtype 与 x 相同
• eps: float，默认 1e-6

语义（必须逐字实现）：

var = mean(x^2, dim=-1) # 沿最后一维

rstd = 1 / sqrt(var + eps)

y = x * rstd * w

输出 dtype 与 x 相同。

约束：

• 内部累加必须提升到 float32，禁止用低精度直接累加
• 不允许在 host 侧用 torch 运算代替内核计算（不要 fallback）
• 不要求反向传播，不要求支持非连续输入
• 单文件实现：rmsnorm_triton.py，导出函数 rmsnorm_triton(x, w, eps=1e-6)
• 先不要做 autotune，用固定的 BLOCK 和 num_warps，方便定位问题

```

注意最后一条。让智能体一上来就 @triton.autotune 是个陷阱：编译时间变长、失败原因被掩盖、你还不知道是哪个配置出的问题。先把正确性跑通，调参留到第 6 步。

同时让它写参考实现，这个参考实现是后面所有对拍的基准：

```python

ref.py

import torch

def rmsnorm_ref(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:

"""PyTorch 参考实现：用 float32 计算，最后转回原 dtype。"""

xf = x.to(torch.float32)

var = xf.pow(2).mean(dim=-1, keepdim=True)

rstd = torch.rsqrt(var + eps)

y = xf * rstd * w.to(torch.float32)

return y.to(x.dtype)

```

参考实现本身必须可信。用 float32 甚至 float64 算，别用 half 算——不然两个都不准的东西比出来的结果没有意义。

第 2 步：先提交测试，再让智能体写实现

这一步是整个闭环的秩序来源。顺序是：写参考实现 → 写对拍脚本 → git commit → 才允许智能体改实现文件。

理由很实际：智能体在多次失败后，很自然会去动测试文件、放宽容差、删掉不通过的 case。你 commit 过测试，就能用 git diff 一眼看出来它有没有偷偷改判卷标准。

对拍脚本这样写：

```python

test_correctness.py

import torch

from ref import rmsnorm_ref

from rmsnorm_triton import rmsnorm_triton

TOL = {

torch.float32: dict(rtol=1e-5, atol=1e-6),

torch.float16: dict(rtol=2e-3, atol=1e-3),

torch.bfloat16: dict(rtol=2e-2, atol=1e-2),

}

def max_err(a: torch.Tensor, b: torch.Tensor):

a = a.double()

b = b.double()

diff = (a - b).abs()

denom = b.abs().clamp_min(1e-6)

return diff.max().item(), (diff / denom).max().item()

def one_case(M, N, dtype, scale, seed, eps=1e-6):

g = torch.Generator(device="cuda").manual_seed(seed)

x = (torch.randn(M, N, device="cuda", generator=g) * scale).to(dtype)

w = torch.randn(N, device="cuda", generator=g).to(dtype)

y_ref = rmsnorm_ref(x, w, eps)

y_tri = rmsnorm_triton(x, w, eps)

ae, re = max_err(y_tri, y_ref)

ok = torch.allclose(y_tri.float(), y_ref.float(), **TOL[dtype])

status = "PASS" if ok else "FAIL"

print(f"[{status}] M={M:<6} N={N:<6} dtype={str(dtype):<16} scale={scale:<8} "

f"max_abs={ae:.3e} max_rel={re:.3e}")

return ok

def main():

shapes = [

(1, 1), (1, 7), (3, 127), (5, 128),

(128, 255), (1024, 1023), (4096, 4096), (65536, 128),

]

dtypes = [torch.float32, torch.float16, torch.bfloat16]

scales = [1e-3, 1.0, 1e3]

failed = 0

seed = 0

for M, N in shapes:

for dtype in dtypes:

for scale in scales:

seed += 1

if not one_case(M, N, dtype, scale, seed):

failed += 1

print(f"\n共 {seed} 个用例，失败 {failed} 个")

if failed:

raise SystemExit(1)

if __name__ == "__main__":

main()

```

这里有四个刻意的设计，值得解释：

• shape 覆盖了边界：N=1、N=7（非 2 的幂）、N=127、N=1023 专治 mask 写错和 BLOCK 取整错误；M=65536, N=128 是典型的“行多列少”场景，和 M=4096, N=4096 的行为完全不同。
• dtype 全覆盖：fp16 / bf16 是精度问题的高发区。
• scale 覆盖数量级：1e-3 到 1e3，平方之后再求和，量级差异会被放大。小值容易下溢，大值容易溢出。
• 失败就非零退出：这样它可以挂进 CI。

第 3 步：让智能体生成内核

把规格、参考实现、对拍脚本一起给它，并明确一句：不要修改 ref.py 和 test_correctness.py。

一份能通过上述对拍的实现大概长这样，可以拿来对照它给你的结果：

```python

rmsnorm_triton.py

import torch

import triton

import triton.language as tl

@triton.jit

def _rmsnorm_fwd(X, W, Y, stride_xm, stride_ym, N, eps, BLOCK: tl.constexpr):

row = tl.program_id(0)

cols = tl.arange(0, BLOCK)

mask = cols < N

x = tl.load(X + row * stride_xm + cols, mask=mask, other=0.0).to(tl.float32)

w = tl.load(W + cols, mask=mask, other=0.0).to(tl.float32)

var = tl.sum(x * x, axis=0) / N # 注意除数用真实的 N

rstd = tl.rsqrt(var + eps)

y = x * rstd * w

tl.store(Y + row * stride_ym + cols, y.to(Y.dtype.element_ty), mask=mask)

def rmsnorm_triton(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:

assert x.dim() == 2, "这个版本只支持二维输入"

M, N = x.shape

assert w.shape == (N,)

x = x.contiguous()

y = torch.empty_like(x)

BLOCK = triton.next_power_of_2(N) # tl.arange 要求 2 的幂

num_warps = 4 if BLOCK <= 1024 else 8

_rmsnorm_fwd[(M,)](x, w, y, x.stride(0), y.stride(0), N, eps,

BLOCK=BLOCK, num_warps=num_warps)

return y

```

三个容易翻车的细节，值得逐条盯：

• var = tl.sum(x * x) / N 里的除数是真实的 N，不是 BLOCK。用 BLOCK 会让结果整体偏小，而且只在非 2 的幂的 shape 上暴露——正好是对拍脚本里 N=127 那几条能抓到。
• 两次 tl.load 后面的 .to(tl.float32)。少了它，fp16 输入会在 fp16 里平方求和，误差大到你怀疑人生。
• tl.store 前显式 .to(Y.dtype.element_ty)。不同 Triton 版本对隐式类型转换的宽容度不一样，显式写出来省事。tl.rsqrt 这类 API 在不同版本间可能有变化，以官方文档当前版本为准。

如果 N 会超过单个 block 能承受的范围（比如几万），这套单 block 写法就不够了，需要改成沿 N 方向循环累加两遍（第一遍求平方和，第二遍算输出）。让智能体加循环时，记得让它保留单 block 的快速路径。

第 4 步：跑对拍，让智能体自己修

```bash

python test_correctness.py

```

第一次大概率不会全 PASS。这时候的正确姿势不是你自己去读 kernel，而是把完整输出贴回给智能体，让它自己改：

```text

python test_correctness.py 的完整输出如下。请在只修改 rmsnorm_triton.py 的前提下修复失败用例。

不要改动 ref.py 和 test_correctness.py，不要放宽 TOL，不要删除任何 case。

先解释每个失败用例的根因，再给修改。

```

给它设一个迭代上限，比如 3 轮。3 轮还不过，通常说明规格和实现之间有理解偏差，人工介入比继续对话快。

对于 fp16 / bf16 的用例，偶尔会出现“大部分 PASS，个别 shape 差一点点”的情况。这类要先看清楚误差量级：如果是 1e-2 以上，几乎一定是精度提升没做；如果是 1e-3 附近在容差边缘，才值得讨论容差是否合理。容差调整要有依据，别为了让测试变绿而放宽。

第 5 步：性能基准

正确性过了，才有资格谈性能。基准脚本要输出三个数：延迟、有效带宽、相对基线加速比。

```python

bench.py

import torch

import triton

from triton.testing import do_bench

from rmsnorm_triton import rmsnorm_triton

from ref import rmsnorm_ref

def bytes_moved(M, N, dtype):

"""读 x + 写 y + 读 w，单位 byte。"""

es = torch.tensor([], dtype=dtype).element_size()

return M * N * es * 2 + N * es

def bench(M, N, dtype):

x = torch.randn(M, N, device="cuda", dtype=dtype)

w = torch.randn(N, device="cuda", dtype=dtype)

ms_tri = do_bench(lambda: rmsnorm_triton(x, w))

ms_ref = do_bench(lambda: rmsnorm_ref(x, w))

gbps = bytes_moved(M, N, dtype) / (ms_tri * 1e-3) / 1e9

speedup = ms_ref / ms_tri

print(f"M={M:<7} N={N:<6} {str(dtype):<16} "

f"triton={ms_tri:8.4f} ms ref={ms_ref:8.4f} ms "

f"bandwidth={gbps:8.2f} GB/s speedup={speedup:5.2f}x")

return speedup

if __name__ == "__main__":

for M, N in [(4096, 1024), (8192, 2048), (16384, 4096), (65536, 128), (131072, 512)]:

for dtype in [torch.float16, torch.bfloat16, torch.float32]:

bench(M, N, dtype)

```

几点使用说明：

• do_bench 的参数签名在不同版本里变过（warmup / rep 的单位），以官方文档当前版本为准。它默认会做 warmup 并清理 L2 缓存，比你自己写 time.time() 靠谱得多。
• 有效带宽比延迟更能说明问题。RMSNorm 是典型的访存受限算子，理论下限就是“把数据从显存读一遍、写一遍”。拿着算出来的 GB/s 去对比你这张卡官方规格里的显存带宽，如果只有理论值的很低比例，说明还有优化空间；如果已经很接近，那再怎么改 kernel 也就这样了。
• 别只看一个 shape。小 M 时启动开销占主导，加速比会很难看，这是正常的，不代表 kernel 写得差。要把 shape 曲线拉出来看趋势。
• 基线要公平。用第 1 步那份 float32 的参考实现当基线是公平的——它同样做了精度提升。如果你的 PyTorch 版本自带 RMSNorm 模块，也可以拿它当基线，但要知道它内部实现可能和你的不是一回事。所有对比都要在同一张卡、同一个进程里跑。

跑之前先确认机器上没有别的进程在占 GPU：

```bash

nvidia-smi

```

第 6 步：判定留还是弃

结果出来了，接下来是决策，不是继续折腾。给一张可以写进 PR 描述的判定表：

“回炉”的动作按成本从低到高排：

1. 调 num_warps（常见是 4 / 8），扫一遍；

2. 调 BLOCK 策略，比如让一个 program 处理多行，减少 program 数量；

3. 上 @triton.autotune，给定几组配置让它自己挑——注意 autotune 的 key 要包含 N 和 dtype，否则会拿错配置；

4. 融合上下游算子（比如把 residual add 并进来），减少一次显存往返。

如果第 1、2 步扫完还是打平基线，那就诚实地弃掉。“打平”意味着这套复杂度不值得：多维护一个手写 kernel、多一份编译依赖、多一个在换卡换驱动时可能出问题的地方，换不来收益。把它记进决策记录里，写清楚测了哪些 shape、什么条件下不划算，下次有人想重做的时候能省一轮。

常见坑与排错

越界读，但结果“看起来对”。忘了 mask 或者 other 没设，越界读到垃圾数据。小 shape 时垃圾恰好是 0，测试就绿了。对策是对拍脚本里必须有 N=1、N=127 这类用例。

BLOCK 取了非 2 的幂。tl.arange 要求长度是 2 的幂，会直接编译报错。用 triton.next_power_of_2(N)。

除数是 BLOCK 不是 N。前面提过，会在非 2 的幂的 shape 上系统性偏小。

没做 float32 提升。fp16 输入下误差可能到 1e-2 量级。所有 tl.load 之后都 .to(tl.float32)。

参考实现自己就是错的。最隐蔽的一类。写参考实现时只用 PyTorch 高层算子（mean、rsqrt、逐元素乘），别在里面手写 for 循环。

智能体偷偷 fallback。检查生成的 rmsnorm_triton.py 里有没有 x.pow(2).mean(...) 这种句子。有就是它放弃了写 kernel。可以直接加一条断言式检查：内核里出现 torch.mean、torch.sum 就等于失败。

智能体偷偷改测试。git diff test_correctness.py 每次必看。

基准只跑了一次。GPU 时钟、其他进程、编译缓存都会影响。用 do_bench，并且同一配置跑两三次看稳定性。

把首次运行时间算进基准。第一次调用包含 JIT 编译，可能几百毫秒。warmup 是必须的。

在 CPU 上跑。Triton 内核需要设备端执行，x 忘了 .cuda() 会报一堆看不懂的错。

版本 API 差异。tl.rsqrt、do_bench 的参数、tl.store 的隐式转换行为在不同版本间都变过。报错先去查官方文档当前版本的 API，别照抄旧代码。

下一步建议

闭环跑通之后，往这几个方向扩展：

• 把对拍挂进 CI。只跑小 shape（比如 M<=256, N<=1024）的子集，几十秒能跑完，能挡住绝大多数回归。大 shape 和性能基准留在本地按需跑。
• 加反向传播。用 torch.autograd.Function 包一层，写 backward 的 kernel，然后拿 torch.autograd.gradcheck 做对拍。反推公式比前向容易写错，对拍的容差通常要放宽到 1e-4 量级（float64 模式下）。
• 做多 shape 的 autotune。把第 5 步扫出来的 shape 曲线变成 autotune 的 key，让小 shape 和大 shape 走不同配置。
• 拿 torch.compile 当参照系。它能生成 Triton 代码，看看它给你的算子生成了什么，往往能发现你没想到的写法。
• 把判定表沉淀成模板。下次再来一个融合算子需求，规格、对拍、基准、决策这四段直接复用，你要写的只剩规格里那几行语义描述。

整条流水线的价值在于：它把“这个 kernel 能不能用”从主观判断变成了三个可复现的数字——对拍通过率、有效带宽、加速比。智能体负责写代码，你负责定义什么叫“通过”。