动画图解 LLM 系统 · RMSNorm

RMSNorm 与 GPU 算子优化

一行公式，一个 kernel：从数学推导到跑满 HBM 带宽

按 → 或点击推进动画 · 页面底部的进度条可以拖动回看

目录

目录

  1. RMSNorm 的公式与几何意义
  2. 和 LayerNorm 的区别：为什么可以不减均值
  3. 反向传播：\(\partial L/\partial x\) 与 \(\partial L/\partial g\)
  4. 数值精度：为什么一定要 FP32 累加
  5. memory-bound：roofline 上的位置
  6. eager PyTorch 的 8 个 kernel
  7. 融合 kernel：行到 block、128-bit 读、warp shuffle、shared memory、寄存器缓存
  8. 与残差相加、FP8 量化的融合
  9. Triton 实现
  10. 反向 kernel 与 dg 的两阶段归约
  11. 实际能达到的带宽
形状 · 从张量到一行

归一化 hidden 维，gamma 在 token 间共享

X ∈ R^(R×D)，γ ∈ R^D，Y ∈ R^(R×D)
R = batch × sequence，D = hidden size

前向按行归约

每个 token 独立计算均方值。γ 是每个通道的可训练增益，广播到所有行；它不是每行一个数。

反向有两个方向

dx 按行归约，因为同一行共享均方根；dγ 沿 token 维求和，因为所有 token 使用同一组 γ。

默认只归一化最后一个维度。多维 normalized_shape 的接口需要先明确归约轴，不能直接沿用一行一个 hidden vector 的地址公式。

公式

RMSNorm：只按均方根缩放

$$y = \frac{x}{\sqrt{\tfrac{1}{d}\sum_{j=1}^{d} x_j^2 + \varepsilon}} \odot g$$

输入是一个 token 的 hidden state：一行 \(d\) 个数（这里画 16 个，数值示意），大小参差不齐。

先逐个平方，再求平均：这就是 mean(x²)，均方。

开根号得到 RMS（均方根），它刻画了这一行"整体有多大"。

每个数除以 RMS：忽略很小的 ε 时，这一行的 RMS 被拉回 1，方向不变，只改长度。\(\varepsilon\)（常取 1e-5 或 1e-6）防止除零。

最后逐元素乘可学习的增益 \(g\)（长度 \(d\)），让每个通道自己决定尺度。没有偏置 \(b\)。

演算 · 让每个数走一遍

平方、求均值、取倒数，再缩放

取 x = [1, 2, 3, 4]，γ = [1, 1, 1, 1]。为方便手算，这一页令 ε = 0；实际实现用模型配置里的 ε。

平方和 30，均方值 7.5，r = 1/√7.5 ≈ 0.365148。最终 y ≈ [0.3651, 0.7303, 1.0954, 1.4606]。

对比

LayerNorm 做两件事，RMSNorm 只做一件

同一组数（示意）分别送进 LayerNorm 和 RMSNorm。

LayerNorm 先算均值 \(\mu\)；RMSNorm 只算 RMS，不需要均值。

LayerNorm 减掉均值，把这一行挪到 0 附近（re-centering）；RMSNorm 跳过这一步。

两边都做缩放（re-scaling）：LayerNorm 除以标准差，RMSNorm 除以 RMS。均值恰好为 0 时，两者完全相同。

少一次归约、少一个偏置，效果却与 LayerNorm 相当，运行时间少 7%–64%；LLaMA、Qwen、DeepSeek 等模型都用 RMSNorm。

不变性

保留缩放不变性，放弃平移不变性

左边是输入，右边两栏分别是 LayerNorm 与 RMSNorm 的输出（数值示意）。

把输入整体乘 2：分子分母同时翻倍，两种输出都不变。这是缩放不变性：权重或激活整体变大，不会改变归一化后的结果。

把输入整体加 1.5：LayerNorm 减均值时把它抵消了；RMSNorm 的输出跟着变。

实验表明，归一化稳定训练的作用主要来自缩放不变性；平移由前后的线性层自己就能吸收。所以去掉均值几乎不损失效果，却省下一次归约。

位置

Pre-norm Transformer 里的 RMSNorm

残差流（下方横线）一路往前；每个子层先对它做 RMSNorm，再把结果加回残差。

每层两次（attention 前、MLP 前），最后再加一次：Llama 3 8B 有 32 层，一次前向调用 65 次。

Qwen3 这类模型还在每个头的 q、k 上做 RMSNorm（QK-Norm），每层再多两次。

每一次调用都要把整个 [tokens × hidden] 读一遍、写一遍；算子本身很简单，成本几乎全在搬数据。

反向传播

梯度沿计算图流回来

前向：x 一路直接走到乘法，另一路经过平方、求均值、rsqrt 得到 \(r\)；\(\hat x = x\,r\)，\(y = \hat x \odot g\)。

梯度 \(dy\) 先到 \(\times g\)：一份给 \(g\)，要对所有 token 求和；一份给 \(\hat x\)。

\(d\hat x\) 再分两路：直接一路乘 \(r\)；另一路经过 \(r\)——\(r\) 依赖整行，所以要先做一次行内归约 \(\sum_j d\hat x_j x_j\)。

两路在 \(x\) 处相加。前向一次行内归约，反向也是一次行内归约，外加一个跨行的 dg 归约。

反向传播

推导

记 \(r = \big(\tfrac1d\sum_j x_j^2+\varepsilon\big)^{-1/2}\)，\(\hat x_i = x_i r\)，\(y_i = g_i \hat x_i\)。先求 \(r\) 对 \(x_j\) 的导数：

$$\frac{\partial r}{\partial x_j} = -\tfrac12\Big(\tfrac1d\textstyle\sum x^2+\varepsilon\Big)^{-3/2}\cdot\frac{2x_j}{d} = -\frac{r^3}{d}\,x_j$$

链式法则，记 \(\hat g = dy \odot g\)：

$$dx_j = \sum_i \hat g_i\Big(\delta_{ij}\,r + x_i\frac{\partial r}{\partial x_j}\Big) = r\,\hat g_j - \frac{r^3}{d}\,x_j\sum_i \hat g_i x_i = r\Big(\hat g_j - \hat x_j\cdot\tfrac1d\textstyle\sum_i \hat g_i\hat x_i\Big)$$

$$dg_i = \sum_{\text{tokens}} dy_i\,\hat x_i$$

前向把每行的 \(r\)（rstd，一个 FP32）存下来，反向就不必重算那次归约；反向需要的输入是 \(x,\ dy,\ g,\ r\)。

演算 · 两条梯度路径相加

归约项会修正整行的梯度

x = [1, 2]，γ = [1, 2]，dy = [3, −1]，ε = 0。r = 1/√2.5，u = dy⊙γ = [3, −2]。

a = mean(u⊙x) = −0.5；dx = r u − r³ x a ≈ [2.023858, −1.011929]。dγ = dy⊙x r ≈ [1.897367, −1.264911]。

数值 · ε 也属于模型定义

小输入时，ε 会明显改变结果

x = [10⁻⁴, 10⁻⁴]，γ = 1，ε = 10⁻⁵
mean(x²) = 10⁻⁸；y ≈ [0.031607, 0.031607]

把输入乘 10，均方值变成 10⁻⁶，输出约 [0.301511, 0.301511]，并没有保持不变。ε 主导时，RMSNorm 更接近一次线性缩放。

写参考实现时

ε 加在均方值里面：rsqrt(mean(x²) + ε)。移到 sqrt 外面、或改用另一套默认值，都会变成不同运算。

检查零输入

x = 0 时 y = 0，但 dx = dy⊙γ/√ε 一般不为零。用“输出为零所以梯度为零”会漏掉真正的导数。

数值精度

平方和必须用 FP32 累加

往累加器里连加 4096 个 1.0：FP32 得到 4096，BF16 卡在 256 不动了。

BF16 只有 8 位有效数字：在 [256, 512) 里相邻两个可表示数相差 2。256 + 1 = 257 落在正中间，就近取偶回到 256。

真实场景：4096 个标准正态数的平方和约 4140，BF16 顺序累加只得到 2048（一次模拟）。RMS 被低估 √2 倍，输出整体放大 1.42 倍。

FP16 还有溢出：激活里出现 300 这样的值并不罕见，300² = 90000 超过 FP16 上限 65504，直接变成 inf。所以 kernel 一律读入后转 FP32、在 FP32 里累加和缩放，写回时再转回 BF16。

GPU

H100：算力上限要看用哪种计算单元

H100 SXM 有 132 个 SM，FP32 CUDA core 峰值约 67 TFLOPS，HBM3 带宽 3.35 TB/s：两者相除约 20 FLOP/字节。RMSNorm 的平方、归约和缩放主要走这条 FP32 路径。

每个 SM 有寄存器与 shared memory，32 个线程组成一个 warp。大矩阵乘走 Tensor Core，BF16 稠密峰值可达 989 TFLOPS；评估归约 kernel，则看 FP32 CUDA core 的算力上限。

Roofline

RMSNorm 在屋顶的斜坡最底下

Roofline：横轴是算术强度（每字节做多少 FLOP），纵轴是能达到的算力。拐点左边受带宽限制，右边受算力限制。

RMSNorm 每个元素约 4 次运算（平方、累加、乘 r、乘 g），却要读 2 字节、写 2 字节：强度约 1，位于 FP32 拐点 20 左侧。图中屋顶采用 FP32 CUDA core 峰值；Tensor Core 矩阵乘有另一条更高的屋顶。

算一笔账：8192 个 token、d = 4096 时，按峰值估算，搬数据下限约 40 µs，FP32 运算下限约 2 µs。优化 RMSNorm 只有两件事：少搬字节，以及让搬运接近峰值带宽。

带宽账本 · 写清读写假设

4096 × 8192 的矩阵，至少要搬多少？

输入、输出与 γ 为 BF16；γ 足够小且在缓存里复用；每行 r 用 FP32 写出。假设一次读 X、一次写 Y，无 spill。

X + Y = 4RD = 134,217,728 B = 128 MiB
r = 4R = 16 KiB；γ 冷读约 16 KiB
计算项结果
按约 4RD FLOP 粗估约 134.2 MFLOP
FP32 67 TFLOP/s 计算下界约 2.00 μs
HBM 3.35 TB/s 传输下界约 40.1 μs

这不是预测实测延迟；launch、归约、缓存、访存效率都可能抬高耗时。若耗时 50 μs，按 128 MiB 口径的有效带宽约 2.68 TB/s。

Eager PyTorch

逐行执行：8 个 kernel，数据来回跑

# Hugging Face LlamaRMSNorm.forward，拆开写
x32  = x.to(torch.float32)        # R2 W4
sq   = x32.pow(2)                 # R4 W4
var  = sq.mean(-1, keepdim=True)  # R4
var  = var + eps                  # 每行 1 个数
r    = torch.rsqrt(var)           # 每行 1 个数
xn   = x32 * r                    # R4 W4
xn16 = xn.to(torch.bfloat16)      # R4 W2
y    = weight * xn16              # R2 W2

# 注释：每个元素读/写的字节数

eager 模式下每个算子是一个独立的 kernel：从 HBM 读输入，算完写回 HBM，下一个 kernel 再读回来。

中间结果都是 FP32，体积是输入的两倍；方差、rsqrt 这些每行一个数的小 kernel 几乎没活，却各占一次 launch。

合计 8 次 launch，每个元素来回搬 36 字节；而这个算子本质上只需要读 x、写 y：4 字节。

融合

一个 kernel：读一次，写一次

融合 kernel 把整行读进 SM，平方累加、归约、乘 \(r\) 和 \(g\) 全在片上完成，中间结果不落 HBM。

每个元素 4 字节（g 只有 d 个数，被所有行共享，几乎全在 L2 命中），流量是 eager 的 1/9，launch 从 8 次变成 1 次。torch.compile 也能生成这样的融合 kernel。

并行划分

一行交给一个 thread block

输入是 [M, d]：M 个 token，每个 d 维。行与行之间完全独立，所以 grid 就开 M 个 block，一个 block 负责一行，由硬件分派到各个 SM。

block 内部：比如 256 个线程 = 8 个 warp，d = 4096 时每个 warp 分到 512 个元素，每个线程 16 个。

d 较小（比如 ≤ 1024）时，一个 warp 处理一行就够了：归约只在 warp 内做，连 shared memory 都不需要。

向量化读取

128-bit 读：一条指令搬 8 个 bf16

标量读：每个 lane 每条指令拿 1 个 bf16（2 字节）。32 个 lane 相邻，访问是合并的，但 512 字节要发 8 条指令。

向量读（uint4 / LDG.E.128）：每个 lane 一次拿连续 16 字节 = 8 个 bf16，一条指令就搬完，访存指令数降到 1/8。

整行 4096 个元素就是 16 次这样的 warp 读：8 个 warp 各读 2 次，每个线程手里 16 个数。要求行首地址 16 字节对齐、d 是 8 的倍数。

向量化 · 先满足地址约束

128-bit 读写也要处理尾部与对齐

合法的快速路径

每次读 8 个 BF16，对应 16 字节。基址和行起点需满足相应对齐，最后一维连续；按实际 stride 计算每行起始地址。

剩下的元素

D 不是 8 的倍数时，只对完整向量使用宽加载；尾部改用带界限检查的标量/较窄加载，不能越界读后再用 mask 掩盖。

D = 1,025：128 次完整的 8 元素加载 + 1 个尾元素

若用 Triton 填到 BLOCK = 2,048，越界 lane 对平方和贡献 0，但均方值必须除以 D = 1,025，不能除以 BLOCK。

Warp 归约

蝶形交换：__shfl_xor_sync

每个 lane 先把自己那 16 个数的平方和算好（数值示意）。现在要把 32 个部分和加起来。

第 1 轮 offset = 16：lane \(i\) 读 lane \(i \oplus 16\) 的寄存器并相加。数据直接在寄存器间交换，不经过内存。

第 2 轮 offset = 8：每个 lane 已经代表 2 个 lane 的和，现在合成 4 个。

offset = 4、2：每轮翻倍覆盖范围。

第 5 轮 offset = 1 之后，32 个 lane 全都拿到总和：log₂32 = 5 步，for (o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffff, v, o);

Block 归约

跨 warp：经 shared memory 中转

shuffle 只能在 warp 内交换。每个 warp 的 lane 0 把自己的和写进 shared memory 的一个槽位，然后 __syncthreads() 等所有 warp 写完。

warp 0 把 8 个槽位读进自己的 lane，再做一次（更短的）shuffle 归约，得到整行的平方和（数值示意）。

lane 0 算出 \(r = \mathrm{rsqrt}(\Sigma/d + \varepsilon)\)，写回 shared memory，再同步一次，所有线程读到同一个 \(r\)。

寄存器缓存

x 留在寄存器里，只从 HBM 读一次

朴素写法分两遍：第一遍读 x 求平方和，第二遍再读一次 x 去乘 \(r\)。每个元素读 4 字节、写 2 字节。

更好的做法：第一遍读进来的 x 就留在寄存器里，等 \(r\) 算出来直接用。每个元素只读 2 字节、写 2 字节。

代价是寄存器：d = 4096、256 线程时每线程 16 个 FP32，很宽松；d 更大时就得增加线程数、把行暂存在 shared memory，或像 Hopper 上那样用 thread block cluster 把一行分给多个 SM。

寄存器 · 少读一次，也可能产生新代价

每行变宽以后，检查是否 spill

一个 block 的线程要保留自己的 x 分片，还要持有地址、gamma、累加器等状态。线程数固定时，D 翻倍通常会提高每线程寄存器需求。

如果寄存器不够

部分数据 spill 到 local memory，最终又走设备内存路径；本来想避免的第二次读取，可能以更差的访存形式回来。

对比两种策略

A：保留所有 x；B：归约后重新加载 x。B 多读一次，但若显著提高 occupancy、消除 spill，也可能更快。

同时记录 registers/thread、local-memory load/store、活跃 warp 和耗时。只看到源代码“读一次”，并不能证明机器码只访问一次 HBM。

CUDA

把这几步拼成一个 kernel

__global__ void rmsnorm_fwd(const bf16* x, const bf16* g, bf16* y,
                            float* rstd, int d, float eps) {
  const int row = blockIdx.x;                 // 一个 block 一行
  const uint4* xr = (const uint4*)(x + (size_t)row * d);
  float v[ITERS * 8];                         // 这一行在寄存器里的份
  float ss = 0.f;
  for (int i = 0; i < ITERS; ++i) {
    uint4 pk = xr[i * blockDim.x + threadIdx.x];  // 128-bit: 8 个 bf16
    unpack8(pk, v + i * 8);                        // 转 FP32
    for (int k = 0; k < 8; ++k) ss += v[i*8+k] * v[i*8+k];
  }
  ss = block_reduce_sum(ss);                  // shuffle + shared memory
  const float r = rsqrtf(ss / d + eps);
  if (threadIdx.x == 0) rstd[row] = r;        // 留给反向
  for (int i = 0; i < ITERS; ++i)             // 用寄存器里的 x
    store8(y, g, v + i * 8, r, row, i);       // y = bf16(x·r·g)
}

示意代码（省略了尾部处理和类型细节）。ITERS = d / (8 × blockDim.x)，编译期常量，这样 v[] 才能全部放进寄存器。

读：每次 128-bit，转成 FP32 存进寄存器，顺手累加平方。

归约 → rsqrt → 把 \(r\) 存一份给反向，然后用寄存器里的 x 算 \(y\) 写回。整条路径只碰 HBM 两次。

再融合

把残差相加和 FP8 量化也并进来

pre-norm 里 RMSNorm 前面总跟着残差相加：h = h + o; x = RMSNorm(h)。分开做要 10 字节/元素，融合后读 h、o，写回 h 和 x：8 字节。

推理时下一层的 GEMM 若用 FP8，还要再量化一次。融合进来后输出直接写 1 字节的 FP8：13 → 7 字节，3 次 launch 变 1 次。

推理引擎里都有这些融合算子，比如 vLLM 的 fused_add_rms_norm、fused_add_rms_norm_static_fp8_quant、rms_norm_dynamic_per_token_quant，以及按 128 元素分块量化的 rms_norm_per_block_quant。

FP8 量化

动态 per-token 量化：再做一次行内归约

RMSNorm 的一行输出 \(y\)（数值示意），要转成 FP8 E4M3，它能表示的最大值是 448。

动态 per-token 量化需要这一行的 \(\max|y|\)：又一次行内归约（把加法换成 max），整行此时还在寄存器里，几乎不花额外的访存。

\(s = \max|y|/448\)，每个数除以 \(s\) 后舍入到 E4M3 网格；网格是非均匀的，越大的数间隔越宽（[256, 448] 内间隔 32）。\(s\) 每行一个 FP32，随输出一起写回。

静态量化用离线校准好的 \(s\)，省掉这次 max 归约；分块量化则每 128 个元素一个 \(s\)，在 block 内分组做 max。

Triton

同一个 kernel 的 Triton 写法

@triton.jit
def rmsnorm_fwd(X, W, Y, RSTD, stride, N, eps,
                BLOCK: tl.constexpr):
    row = tl.program_id(0)              # 一个 program 一行
    cols = tl.arange(0, BLOCK)          # BLOCK = next_pow2(N)
    mask = cols < N
    x = tl.load(X + row * stride + cols, mask=mask,
                other=0.).to(tl.float32)
    rstd = 1 / tl.sqrt(tl.sum(x * x, axis=0) / N + eps)
    tl.store(RSTD + row, rstd)
    w = tl.load(W + cols, mask=mask).to(tl.float32)
    y = x * rstd * w
    tl.store(Y + row * stride + cols,
             y.to(Y.dtype.element_ty), mask=mask)

# 启动：M 行就是 M 个 program
rmsnorm_fwd[(M,)](x, w, y, rstd, x.stride(0), N, eps,
    BLOCK=triton.next_power_of_2(N), num_warps=8)

Triton 以"一个 program 处理一块数据"来写：program_id 就是行号，tl.arange 一次覆盖整行。N = 5120 时 BLOCK = 8192，多出来的 3072 个位置被 mask 掉。

tl.load 由编译器向量化；tl.sum 被编译成 warp shuffle + shared memory 归约；整行 x 作为一个张量留在寄存器里。

需要手动选的只剩 num_warps 等少数参数。CUDA 版本的那几页优化，Triton 编译器大都替你做了。

反向 kernel

dx：每行再归约一次

反向也是一行一个 program：读 x、dy、g 和前向存下的 \(r\)，算 \(\hat x = x r\)、\(\hat g = dy\odot g\)，行内归约得到 \(c = \mathrm{mean}(\hat x\odot\hat g)\)，写出 \(dx = r(\hat g - \hat x c)\)。

每个元素读 4 字节、写 2 字节。同一个 program 往下处理更多行，\(dy\odot\hat x\) 一直累加在 FP32 寄存器里——这就是 dg 的部分和，最后写回一次。

dg 归约

两阶段：先按 program 求部分和，再按列加起来

dg 是对所有行求和，和行内归约方向正交。让 P 个 program 各自跨步处理行（program p 负责第 p、p+P、p+2P… 行），各自在寄存器里累加。

阶段 1 结束，写出 [P, d] 的 FP32 部分和；第二个小 kernel 按列求和得到 dg。P 一般取 SM 数（H100 上 132），d = 4096 时缓冲区只有 2 MB 左右。

为什么不让每行直接 atomicAdd 到 dg：M 行抢同 d 个地址，冲突严重，而且浮点加法顺序不固定，结果不可复现。另一种做法是让几组行共享一个缓冲区、用锁串行累加，同样是两阶段。

实测

能跑到多少带宽

衡量 memory-bound kernel 用有效带宽：必须搬的字节数 ÷ 耗时。H100 SXM 峰值 3.35 TB/s，实际拷贝也达不到 100%。

精心调优的归约类 kernel（softmax、RMSNorm 等）在 H100 上约 3 TB/s，接近峰值的 90%；编译器生成的版本在超宽的行上会掉到 2 TB/s 左右。eager 实现即使每个 kernel 都跑满带宽，也只有 4/36 的字节是有用的。

交互

token 数少的时候，带宽也跑不满

融合 kernel 耗时
eager 耗时
融合有效带宽
加速比

示意模型：每个 kernel 固定开销 3 µs（launch、启动与收尾），其余时间按 3.0 TB/s 持续带宽搬字节；融合 4 B/元素、1 个 kernel，eager 36 B/元素、8 个 kernel。decode 阶段 M 很小，固定开销占大头，这时把多个算子融合（或用 CUDA Graph）比优化带宽更重要。

正确性 · 先把误差的责任分开

用 FP64 验证导数，再比较低精度 kernel

输入集合要暴露的问题
D = 1、7、1,025归约宽度、padding、尾部加载
全零、极小值、300 与 −300ε、FP16 平方溢出、有限性
非连续行 / 非标准 stride地址计算与接口约束
不同 gamma、多个 token增益广播、dγ 沿 token 求和
固定 dy 的中心差分dx 的归约项、dγ 的维度
数值导数 ≈ [L(x + h eᵢ) − L(x − h eᵢ)] / (2h)

先用 float64、h≈10⁻⁶；低精度差分常被舍入淹没。再以同样的转换与 ε 约定比较 BF16/FP16，分别记录最大绝对误差和相对误差。

测量 · 大矩阵与单 token 分开看

有效带宽只是一个带假设的指标

# x/y: BF16；gamma 在缓存复用，不计其每行重复读取
bytes_estimate = 4 * rows * hidden
for _ in range(20): y = rms_kernel(x, gamma, eps)
a = torch.cuda.Event(enable_timing=True)
b = torch.cuda.Event(enable_timing=True)
a.record()
for _ in range(100): y = rms_kernel(x, gamma, eps)
b.record(); b.synchronize()
ms = a.elapsed_time(b) / 100
print(ms, "ms", bytes_estimate / (ms * 1e9), "TB/s")

rms_kernel 为待测实现；以上需要 CUDA。逐项扫 rows 与 hidden，并记录 warm-cache 或冷缓存策略；数据驻留 L2 时，不能把有效带宽直接当成 HBM 带宽。

排查 · 同一 kernel 的两种慢法

先判断是 launch 不够快，还是数据搬不动

教学测量推理下一步
R=1，D=8192，8 μs只搬约 32 KiB，工作量极少查 launch / 调度；考虑融合邻接算子
R=4096，D=8192，50 μs约 2.68 TB/s，有效带宽较高改算式的收益有限，评估减少流量
同形状扩大 D 后突降可能触及寄存器或 block 资源边界查 spill，比较 block/warp 配置
D=1024 正确，1025 结果偏大可能把归约除数错写为 BLOCK核对真实 D 与 padding mask

这些数字是诊断练习的假设观测。评价融合时，把整个 residual → norm → quantize 路径计时，避免仅看到某个 kernel 变快。

练习 1 · 独立分析

用两维输入手算前向与反向

x = [1, 2]，γ = [1, 2]，dy = [3, −1]，ε = 0。求 y、dx、dγ，并检查 x·dx。
展开参考解答

① r = 1/√2.5 ≈ 0.632456；y = [r, 4r] ≈ [0.632456, 2.529822]。

② u = [3, −2]，a = (3−4)/2 = −0.5；dx = ru−r³xa ≈ [2.023858, −1.011929]。

③ dγ = [3r, −2r] ≈ [1.897367, −1.264911]；x·dx = 0，对应 ε=0 时的正比例缩放不变性。

先独立计算，再展开解答。

练习 2 · 独立分析

padding 的一个除数，会带来多大误差？

x 有 5 个元素且全是 1，γ 全为 1，ε = 0。实现填充到 BLOCK=8，尾部写 0，却把平方和除以 8。输出是多少？正确输出呢？
展开参考解答

① 有效元素的平方和是 5；正确均方值为 5/5 = 1，每个有效输出应为 1。

② 错误均方值为 5/8，r = √(8/5) ≈ 1.264911；所有有效输出被放大约 26.49%。

③ padding mask 只负责让无效 lane 贡献 0；归约分母仍须用真实 hidden size。

先独立计算，再展开解答。

练习 3 · 独立分析

50 μs 的结果还有多少带宽空间？

R=4096、D=8192，BF16 一读一写共 128 MiB。假设 HBM 峰值 3.35 TB/s，实测 50 μs。计算有效带宽和仅按 HBM 下界估计的剩余空间。
展开参考解答

① 128 MiB = 134,217,728 B；带宽 = 134,217,728 / (50×10⁻⁶) ≈ 2.684 TB/s，约峰值的 80.1%。

② 理想下界为 40.1 μs；在相同流量口径下，50/40.1≈1.25×，对应延迟再降约 19.9%。

③ 若数据命中 L2、gamma 被重复读、出现 spill，账本就会变。要先核对流量口径，再解释百分比。

先独立计算，再展开解答。

动手 · 把导数变成可以检查的代码

一次运行，检查 dx 与 dgamma

python3 lab.py --width 7 --eps 1e-5
# 再试：--width 1、--width 13、--eps 0.01

脚本实现前后向，用中心差分验证每个输入和 gamma 的导数，还检查手算结果、全零输入、非 2 次幂的 padding 与 FP32 平方。

默认：dx / dγ 最大差分误差 < 1e−7
通过后打印 PASS

修改题：把 mean(u*x) 改成 sum(u*x)，或让 padding 的归约除以 BLOCK。先预测错误如何随 D 变化，再运行检查。

下载 lab.py · NumPy / CPU
小结

小结

  • RMSNorm = 除以均方根再乘增益；不减均值，通常不加偏置；忽略 ε 时，对输入的正比例缩放保持不变。
  • 反向：\(dx = r(\hat g - \hat x\cdot\mathrm{mean}(\hat g\odot\hat x))\)，\(dg = \sum_{\text{tokens}} dy\odot\hat x\)；前向存下每行的 \(r\)。
  • 平方和一定在 FP32 里累加：BF16 顺序连加 1 会在 256 停住，FP16 平方可能溢出。
  • 大张量时算术强度约 1 FLOP/B，受带宽限制：目标是每元素只读 2 字节、写 2 字节，并接近峰值带宽。
  • 做法：一行一个 block（或 warp）、128-bit 读、warp shuffle + shared memory 归约、x 留在寄存器里。
  • 再往前一步是和残差相加、FP8 量化融合：13 → 7 字节/元素，3 次 launch 变 1 次。
  • 反向的 dg 是跨行归约：P 个 program 各求部分和，第二个 kernel 按列相加。
参考

参考与延伸阅读

系列 · 动画图解 LLM 系统

下一篇：Gradient checkpointing

训练时激活值占掉的显存，能不能用算力换回来

← 回到博客