LayerNorm and Welford Online Statistics
REFERENCE CONTEXT · CUDA 12.2; RTX 2070 SUPER; Turing sm_75
Core Concept
LayerNorm computes row mean and variance before an affine transform. Welford maintains (count, mean, M2) online and supports parallel combination.
Why It Affects Performance
The formula E[x^2]-E[x]^2 can suffer catastrophic cancellation when the mean is large and variance is small. Multi-kernel implementations also reread inputs and write statistics.
Example or Pseudocode
n' = n + 1
delta = x - mean
mean' = mean + delta/n'
M2' = M2 + delta*(x - mean')
variance = M2/nAfter reduction, compute (x-mean)*rsqrt(var+eps)*gamma+beta in the same kernel.
Profiling and Performance Judgment
Track kernel count, DRAM bytes, FP32 statistics, register pressure, and multiple hidden sizes. Compare against torch.nn.functional.layer_norm.
Five-Minute Exercise
Why use FP32 mean and variance for FP16 inputs?
Show answer
Statistics involve long reductions and subtraction of similar values. FP32 reduces overflow, underflow, and cancellation error.