T3 讲解：LayerNorm / RMSNorm 的 forward

① 先看全貌：两个公式

② kernel 逐行讲

drillslayer_norm_kernel（16–38 行）

读和均值

row  = tl.program_id(0).to(tl.int64)
cols = tl.arange(0, BLOCK_N)
mask = cols < N
x    = tl.load(x_ptr + row * x_row_stride + cols, mask=mask, other=0.0).to(tl.float32)
mean = tl.sum(x, axis=0) / N

前四行和 T2 一样。这里 padding 填的是 0（T2 填 -inf），因为下一步是求和：0 不影响 Σx。注意除以的是真实的 N，不是 BLOCK_N。

陷阱 1：padding 减完均值就不是 0 了

drills27–29 行

xc  = tl.where(mask, x - mean, 0.0)
var = tl.sum(xc * xc, axis=0) / N

padding 位置的 x 是 0，但 x − mean = −mean，平方之后是 mean²，会被加进方差。N = 5、BLOCK_N = 8 时：

正确的方差 = (4 + 0 + 0 + 1 + 1) / 5 = 1.2。不用 tl.where 的话，三个 padding 各贡献 16，方差变成 (6 + 48) / 5 = 10.8，错了 9 倍。而且 N 恰好是 2 的幂时没有 padding，测试发现不了这个 bug。

tl.where(cond, a, b) 是逐元素的三元选择：cond 为 True 取 a，否则取 b。Triton 里没有逐元素的 if，都用它。

陷阱 2：方差用哪个公式

参考解用的是两遍法：先求 mean，再求 Σ(x − mean)²。教科书上还有一个“一遍法”：var = E[x²] − E[x]²，只需要同时累加 Σx 和 Σx²。它在数学上等价，在浮点上可能灾难性地出错：

数据的均值远大于它的波动时（LLM 的激活值里经常出现个别很大的通道），两个很大的数相减，有效数字全被抵消掉了，这叫灾难性抵消。两遍法先把均值减掉，再平方，就没有这个问题。

整行在寄存器里，所以“两遍”只是对寄存器里的数据算两次，不需要多读显存。如果行太长、必须分段从显存读，又想只读一遍，就用 Welford 算法：边读边更新均值和 M2，数值上也是稳定的。PyTorch 的 CUDA LayerNorm 就用它（见第 ⑤ 节）。

rstd、仿射变换和写回

drills30–38 行

rstd = 1.0 / tl.sqrt(var + eps)
w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
b = tl.load(b_ptr + cols, mask=mask, other=0.0).to(tl.float32)
y = xc * rstd * w + b
tl.store(y_ptr + row * y_row_stride + cols, y.to(y_ptr.dtype.element_ty), mask=mask)
tl.store(mean_ptr + row, mean)     # 标量写：每个 program 一个数
tl.store(rstd_ptr + row, rstd)

• eps 防止一行全相同时方差为 0、除以 0。LayerNorm 默认 1e-5，RMSNorm 常用 1e-6。
• 存 rstd（1/σ），而不是 σ 或 var：backward 里用到的正是 rstd，乘法比除法便宜。
• 最后两行是标量 store：mean 和 rstd 是标量，mean_ptr + row 是一个地址。
• 为什么要存 mean 和 rstd：backward 计算 dx 需要 x̂ = (x − mean) · rstd。存下来每行只多 8 字节，省得 backward 再做两次归约。

③ RMSNorm：少一次归约

drillsrms_norm_kernel（41–49 行）

x    = tl.load(..., other=0.0).to(tl.float32)
rstd = 1.0 / tl.sqrt(tl.sum(x * x, axis=0) / N + eps)
w    = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
tl.store(y_ptr + ..., (x * rstd * w).to(...), mask=mask)

不减均值，所以 padding 填 0 之后平方还是 0，不需要 tl.where。没有 b，也不存统计量（这题只要求 forward）。

题目页的实测里 RMSNorm 和 LayerNorm 一样快（N = 8192 时都是 176 us）：两者都要读 M×N 的 x、写 M×N 的 y，少一次寄存器里的归约几乎不省时间。memory-bound 的 kernel 只有减少主要数据的字节数才有用。

④ host 端：reshape 和 num_warps

drills52–73 行

x2 = x.reshape(-1, shape[-1])            # [B, S, N] → [B·S, N]
if x2.stride(-1) != 1: x2 = x2.contiguous()
BLOCK_N = triton.next_power_of_2(N)      # 超过 16384 就报错
num_warps = min(max(BLOCK_N // 256, 1), 8)

• reshape 成二维：LayerNorm 只对最后一维归一化，前面所有维度都可以合并成“行”。reshape 在可能时返回视图，不拷贝。
• num_warps 的规则：大约每 256 个元素一个 warp，最少 1、最多 8 个。

• 返回值：y 恢复成原来的形状；mean 和 rstd 始终是 fp32 的 [M]。

⑤ 对照真实实现

• Triton 官方 tutorial 05：triton_layer_norm_fwd_fused 支持任意长的行：用 tritonfor off in range(0, N, BLOCK_SIZE) 分段读，求方差时 triton再读一遍、减均值后平方，也是两遍法，只不过两遍都从显存读。
• FlashInfer 的 CUDA RMSNorm：flashinferRMSNormKernel，Σx² 先在线程内累加，再用 flashinferwarp shuffle 和 flashinfershared memory 两层归约，正是 C2 的结构。
• vLLM 的 RMSNorm：vllmrms_norm_kernel，用 vllmcub::BlockReduce 做 block 归约。F3 会实现它的融合版 fused_add_rms_norm。
• PyTorch 的 CUDA LayerNorm：pytorchRowwiseMomentsCUDAKernel 用 pytorchWelford 算法 一遍算出均值和方差。

⑥ 常见坑与自测

• 忘了 tl.where：N 不是 2 的幂时方差偏大。
• 除以 BLOCK_N 而不是 N：padding 位置被当成了数据。
• 用无偏方差（÷ N−1）：F.layer_norm 用的是有偏方差（÷ N）。
• w、b 不转 fp32：fp16 的 w 和 fp32 的 y 相乘时类型提升规则可能和预期不同，统一转 fp32 最省心。
• mean、rstd 用了输入的 dtype：它们应该是 fp32，backward 需要精度。

源码与资料索引

本仓库：drillssolutions/t3_layernorm.py · drillsmine/t3_layernorm.py · drillstests/test_t3_layernorm.py · drillscontent/t3.md

参考实现：tritonTriton tutorial 05 · flashinferFlashInfer Triton norm · flashinferFlashInfer RMSNormKernel · vllmvLLM rms_norm_kernel · pytorchPyTorch RowwiseMomentsCUDAKernel

API 与背景：tl.where · tl.sqrt · Welford 在线方差算法 · LayerNorm 论文 · RMSNorm 论文