RMSNorm 的两次读取：别把一次归约叫作单遍访存

RMSNorm 省掉了 LayerNorm 的去均值步骤，但这不等于整行只读一次。通常先用输入求平方和，得到归一化系数，再读输入生成输出。如果数据留在缓存或寄存器里，外部流量可能减少；如果输入很长，第二次访问仍需认真计入。

2026 年 9 月的 低精度 FA4 实现说明讨论 RMSNorm 与量化生产者融合。值得追的是中间结果如何产生和复用，而不是只把循环数量从二改成一。

先算逻辑字节账

对 D 个 FP32 元素，输入和输出各占 4D 字节。即使只算一次读与一次写，也至少有 8D 字节逻辑载荷。两次读取输入，再读 gamma、写输出，是 16D；D=4096 时为 65536 字节。这个数没有包括权重缓存复用，也不是 DRAM 性能计数器结果。

平方和归约与输出缩放是两个阶段，输入复用需要实际存储承载。

import math

def rmsnorm(x,gamma,eps=1e-5):
    if not x or len(x)!=len(gamma) or eps<=0:
        raise ValueError("invalid normalization input")
    if not all(math.isfinite(v) for v in x+gamma):
        raise ValueError("finite values required")
    scale=1/math.sqrt(math.fsum(v*v for v in x)/len(x)+eps)
    return [v*scale*g for v,g in zip(x,gamma)]

x=[-2.,0.,1.,3.];g=[1.]*4
y=rmsnorm(x,g)
assert math.isclose(sum(v*v for v in y)/4,3.5/(3.5+1e-5))
assert rmsnorm([0.,0.],[1.,1.])==[0.,0.]
assert 4*4096*4==65536
print("RMSNorm: formula and logical byte count passed")

Welford 解决的问题不同

LayerNorm 需要均值和方差，Welford 提供一种增量更新统计量的方法，有助于避免直接做 E[x²]-E[x]² 的抵消问题。它没有自动替后面的输出归一化保留全部输入。把统计量计算称为单遍是合理的，把整行算子称为单遍输入访问则需要说明数据保存在哪里。

融合后续量化，可以避免先写 FP32 中间张量再读回来，但要同时管理输出尺度、舍入与异常值。寄存器数量、共享内存和行宽会影响能否保留数据；占用更多寄存器也可能降低并行度。

小工作集与大工作集的两种解释

同样是两次逻辑读取，数据可能来自不同层级。若一行数据在第一次归约后仍留在高速缓存，第二次读取未必再次访问 DRAM；若行很长或同时有许多行竞争缓存，就不能作这种假设。字节推导描述程序请求了多少数据，计数器和运行轨迹才回答这些请求实际去了哪里。

gamma 也可能被多行复用。把每行读取 gamma 的逻辑量都算进账本没有错，但估计外部带宽时应考虑缓存复用。反过来，若 gamma 的布局导致不连续访问，实际传输可能大于有效元素载荷。理论最低流量、逻辑流量和实测流量应各有一栏，不能用最小的一栏解释所有场景。

再看融合量化。若归一化输出需要按整行最大幅值重新决定尺度，就又引入了一次统计依赖。一个实现可能保留整行数据，另一个实现可能多读一次，或者使用校准尺度避免动态统计。三者的工作区、误差和调度条件不同，不宜只比较最终输出都是 FP8。

数值测试也应有区分力。全零输入可以检查 epsilon 路径，包含不同量级的输入可以观察归约误差，非单位 gamma 可以暴露错误的缩放顺序。把所有 gamma 都设成一，虽然容易核算，却无法覆盖参数索引。扩大实现范围之前，应为这些条件分别增加参考检查。

批量行之间也需要独立统计量。若把相邻行的平方和合并，再统一缩放，会改变归一化轴。只测试重复的相同行数据，可能看不出这种错误；让相邻行的幅值和 gamma 各不相同，才更容易验证行索引。行内公式成立和批次索引成立，应分别取得证据。

本例是标量数值与字节推导检查，没有验证极大浮点数的溢出鲁棒性，没有测 SIMD、GPU、内存带宽或真实模型。性能表应由目标设备测量产生；在此之前，把“减少了哪一次物化”说清，比写一个未经验证的倍数有用。