Meta's FlashNormAttention Hides 90% of Transformer Normalization Costs

Normalization layers are everywhere in modern deep learning -- every transformer block has at least one, and Meta's largest recommendation models are saturated with them. But they come with a hidden tax: in the Kunlun architecture deployed on Meta's largest RecSys training foundation model, normalization takes up roughly 20% of the total training latency. That's 20% of your GPU budget spent on operations that don't use Tensor Cores at all. Meta's kernel team just published a detailed breakdown of how they're clawing that back.

The full blog post introduces three novel techniques -- Lazy Pre-Norm, Multi-CTA Norm Fusion, and FlashNormAttention -- that together can hide up to 90% of normalization latency by fusing it directly into adjacent GEMM and Attention kernels. The code is open-sourced under the facebookresearch/ads_model_kernel_library repository.

Why normalization is so hard to fuse

The core problem is a tiling mismatch. On a GPU, work is divided into tiles -- small rectangular chunks of a matrix that fit in shared memory and get processed together. Normalization is by nature a reduction operation that requires access to data along an entire dimension. For LayerNorm and RMSNorm, a typical kernel tiles the input along the outer dimension but not the inner, meaning each CTA always needs to load entire rows of data. By comparison, a typical GEMM is tiled in both dimensions, meaning each tile does not span an entire row, making a following row-wise normalization impossible.

The naive fix -- stretching the GEMM tile to span the full row -- runs into two walls: it degrades GEMM performance by forcing a suboptimal tile shape, and it hard-limits how large the embedding dimension N can be. On a Blackwell GPU with 228KB shared memory, this restricts N to at most 512 for the kernel to even be able to run. That's fine for small shapes, but modern LLMs and recommendation models routinely exceed it.

Three techniques, one goal

Meta's team developed three distinct strategies to work around this, each targeting a different part of the model graph: