Novel GPU kernel fusion techniques are presented for LayerNorm and RMSNorm operations, which typically consume 10-20% of training latency in LLMs and recommendation models. Three main contributions are covered: (1) naive epilogue fusion for small N dimensions yielding 17-32% latency savings; (2) Lazy Pre-Norm, a prologue fusion technique that exploits the commutativity of row-wise multiplication to delay RMSNorm's elementwise step until after the GEMM k-loop, hiding most norm computation behind TensorCore ops; (3) Multi-CTA Norm Fusion using CTA clusters and Distributed Shared Memory to collaboratively compute reductions across CTAs, enabling epilogue fusion for larger N. A 'fusion regrouping' trick ensures backward passes also achieve epilogue fusion by fusing norms with different linear layers than in forward. Finally, FlashNormAttention fuses both a pre-LayerNorm and post-RMSNorm into a GDPA attention kernel, achieving up to 35% kernel speedup, with detailed optimizations for memory pressure (SMEM/TMEM reuse, register subtiling) and pipeline stalls (warp specialization, register pre-loading). Benchmarks are on NVIDIA B200 GPUs using bfloat16.
