source: pytorch blog: towards free normalization: fusing normalization into gemm and attention kernels
level: technical
normalization layers like layer norm and rms norm are common in large language models and recommendation systems, but they are memory-bound and waste compute. in meta's kunlun architecture, normalization can take 20% of training time. the main challenge is that normalization needs whole rows of data, while typical matrix multiplication tiles data in both dimensions, causing a mismatch. a naive fusion forces the matrix multiplication tile size to span the entire inner dimension, which works for small sizes but fails for larger ones due to shared memory limits and poor tiling.
lazy pre-norm tackles pre-rms norm fusion by delaying the elementwise multiplication until after the matrix multiplication. it computes the row-wise scaling factor in parallel with the matrix multiplication loop, then applies it to the output. this avoids the cyclic dependency where the scaling factor is needed before the loop finishes. the technique works only for rms norm without elementwise affine parameters and cannot handle layer norm. it still repeats some computation across thread blocks, but the overlap with tensor core operations makes it efficient.
multi-cta norm fusion addresses post-norm fusion by using cta clusters and distributed shared memory. multiple cooperative thread arrays work on the same rows, splitting the inner dimension and communicating via shared memory to compute normalization. this allows fusing post-layer norm or rms norm with matrix multiplications without forcing a single tile to cover the whole row. the approach generalizes to attention kernels with flashnormattention, which fuses both a layer norm and an rms norm around an attention operation, achieving up to 35% speedup.
why it matters: reducing normalization overhead frees up gpu compute for other tasks, directly speeding up training and inference of large models.
source: pytorch blog: towards free normalization: fusing normalization into gemm and attention kernels