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:

• Lazy Pre-Norm -- fuses RMSNorm that precedes a linear layer (pre-norm) by exploiting a mathematical identity: since RMSNorm without elementwise affines is just a row-wise multiplication, you can commute it past the matrix multiply. The key observation is that (A * rstd[:, None]) @ B = (A @ B) * rstd[:, None]. This means the elementwise computation can be lazily computed and delayed until after the whole k-loop is done, effectively becoming an epilogue. The reduction (computing the normalization scale factor) runs in parallel with the Tensor Core matmul, and the elementwise rescaling happens at the very end -- no cyclic dependency.
• Multi-CTA Norm Fusion -- handles post-norm (normalization after a linear layer) for larger N values, where the naive approach breaks down. The technique borrows an idea from Quack and extends it beyond standalone norm kernels to fused kernels. Quack norm kernels leverage CTA clusters to partition large N among different CTAs in the same cluster, and let them collaborate on a single reduction across N by communicating necessary data with each other via distributed shared memory. This avoids the cost of going back to global memory (HBM) for the cross-CTA communication. The practical limit on Blackwell is N up to 4096 (512 per CTA × max cluster size of 8).
• FlashNormAttention -- applies both of the above ideas simultaneously to a full attention block, fusing a LayerNorm pre-norm, an RMSNorm post-norm, and two residual connections directly into a single GDPA (Generalized Dot-Product Attention) kernel. This is like a megakernel performing all operations in a module, but it differs in that it aims not just to save kernel launch costs, but to save the total amount of data transfer to/from HBM.

The engineering depth behind the numbers

The FlashNormAttention kernel is where the real complexity lives. Fusing two norms and two residual connections into a single attention kernel creates severe pressure on shared memory and registers. The team applied several hardware-level tricks to make it work:

• SMEM/TMEM buffer reuse: Non-overlapping tensors share the same memory buffers. For example, the output buffer is reused to temporarily hold intermediate values once they're no longer needed.
• TensorCore accumulate: Instead of keeping the LayerNorm output in shared memory and reading it out after the matmul, the team leverages the MMA semantic supported by tcgen05 in TMEM, directly keeping the value in the TMEM buffer allocated for the matmul and offloading the addition to TensorCore.
• Register subtiling: To avoid register spilling (where registers overflow to slower memory), tensors are cut into chunks and loaded one piece at a time for normalization computation.
• Warp specialization: The kernel uses five specialized warp partitions -- load, MMA, activation, epilogue, and a fifth partition dedicated to prologue LayerNorm computation -- in order to better overlap it with TensorCore as well as other CUDA Core operations.

The backward pass adds another layer of difficulty. Epilogue fusion in the forward pass naturally becomes prologue fusion in backward -- which is exactly the bad case (norm on the critical path, blocking every matmul iteration). The solution is elegant: fuse the normalization op with different linears in forward vs. in backward, resulting in efficient epilogue fusion in both directions.

What the numbers actually look like

Benchmarks were run on NVIDIA B200 GPUs in bfloat16, on real-world shapes from Meta's ads recommendation traffic:

• Lazy Pre-Norm hides between 41% and 98% of RMSNorm kernel latency across common matrix shapes.
• Multi-CTA Norm Fusion hides significant LayerNorm and RMSNorm latency for epilogue fusion across (K, N) configurations up to 2048.
• FlashNormAttention achieves up to 35% kernel speedup for full Attention blocks with pre-norm, post-norm, and residual connections.

The 90% figure refers to the best case for Lazy Pre-Norm on small shapes -- essentially making normalization free by fully hiding it behind Tensor Core execution. The 35% end-to-end attention block speedup is the more practically relevant number for most architectures.

The tooling stack

Two kernel DSLs were used to build these kernels. Helion is a new kernel authoring DSL tightly integrated with PyTorch 2 that compiles down to Triton, raising the abstraction level to simplify kernel development, minimize boilerplate, and significantly enhance maintainability and portability. It was used for the Lazy Pre-Norm kernels where exhaustive autotuning over non-standard tile shapes was critical. TLX, a set of lower-level Triton DSL extensions with hardware-aware GPU execution control, was used for the Multi-CTA kernels where fine-grained control over CTA cluster scheduling was essential.

What this means for the field

One of the key challenges inherent to the Transformer architecture is the requirement to support numerous non-linear transformations that involve normalization. Each decoder block typically contains at least one Softmax operation and two LayerNorms. The computation of the corresponding normalization scaling factors becomes a major bottleneck because it requires spatial collective operations. This work directly attacks that bottleneck at the kernel level rather than at the architecture level.

The practical implication is that the common assumption -- that normalization is a fixed overhead you simply pay -- needs updating. These techniques show that for models with large embedding dimensions and frequent normalization (which describes most production LLMs and recommendation systems), a significant fraction of that cost can be recovered without changing the model architecture at all. The constraint is hardware: the Multi-CTA approach currently requires Hopper or Blackwell GPUs with CTA cluster support, and N is bounded at 4096 on Blackwell's portable cluster size of 8. For teams running on older hardware or with very large embedding dimensions, the gains will be more limited.

The code for both the Multi-CTA Norm Fusion and FlashNormAttention megakernel is publicly available, making this directly usable for teams building on similar architectures.