Towards Free Normalization: Fusing Normalization into GEMM and Attention Kernels – PyTorch

Search

Close Search

Blog

Towards Free Normalization: Fusing Normalization into GEMM and Attention Kernels

By Jacky (Junqing) Zhou, Hongtao Yu, Jackie (Jiaqi) Xu, Menglu Yu, Ethan Che, Han Xu, Darren Liu, Peng Chen (Dev Infra), Daohang Shi, Max LeungJuly 10, 2026No Comments

Featured projects

Code available at: https://github.com/facebookresearch/ads_model_kernel_library/tree/main/multi_cta_norm_fusion and https://github.com/facebookresearch/ads_model_kernel_library/tree/main/gdpa_megakernel

TL;DR

In this blog post, we present various novel kernel fusion techniques for common normalization ops like LayerNorm and RMSNorm, which provide significant speedup by reducing the memory-IO overhead of these highly memory-bound kernels. We start with a brief overview of the modeling importance as well as performance challenges of normalization ops common in both LLMs and ads recommendation models, then present the novel strategies in tackling the performance bottlenecks, including Lazy Pre-Norm and Multi-CTA Norm Fusion. We show that such techniques can hide as much as 90% of a normalization kernel’s latency by fusing with GEMMs. In the end, we present the FlashNormAttention algorithm where we apply fusion for multiple normalizations around an attention kernel like GDPA [1], achieving up to 35% kernel speedup.

This work is done with primarily two kernel DSLs: TLX, a set of Triton DSL extensions with lower-level, hardware-aware support for GPU execution control; and Helion, a high-level DSL that excels at developer velocity, portability, and comprehensive autotuning. Benchmarks are performed with data type bfloat16, on NVIDIA B200 GPUs in Meta’s data centers with a 750 W power cap.

Introduction

Normalization techniques have become indispensable in most deep learning architectures due to their excellent effectiveness in stabilizing training and accelerating convergence. In particular, traditional normalization across the innermost embedding dimension (e.g. LayerNorm, RMSNorm) has been the most common and ubiquitous type in modern Large Language Models as well as recsys models like Meta’s ads models. For example, in the Kunlun [2] architecture deployed on Meta’s largest Recsys training foundation model, the Generative Ads Model (GEM) [3], LayerNorm/RMSNorm exists in nearly all key components, such as Multi-Head Attention, Hierarchical Seed Pooling, and GDPA-enhanced [1] PFFN.

However, the ubiquity of normalization also brings a difficult performance challenge: it’s highly memory-bound with no TensorCore utilization. This hinders us from saturating hardware compute capabilities in model training. Using Kunlun [2] as an example, normalization takes up roughly 20% of the total training latency there. This means we immediately lose 20% of our hardware’s compute throughput without optimization. In a typical LLM which is more compute-bound, normalization could still take roughly 10% of total latency.

To address this, we must design our normalization kernels in an IO-aware way, carefully saving memory IO costs without compromising computation accuracy via kernel fusion, and also overlapping these memory-/CUDACore-heavy ops with TensorCore-intensive ops. Since most normalization operations follow or precede matmul operations (e.g. MLP, Attention), our work focuses on how to efficiently fuse norms with matmuls. We start by describing multiple strategies in efficiently fusing norms with single GEMMs, and in the end present the FlashNormAttention algorithm, fusing both a LayerNorm and an RMSNorm into attention.

Note : In the following benchmark results, we disable elementwise affines unless otherwise specified, as we find those to incur significant performance overhead with marginal model quality effects in our models. Also, unless otherwise specified, the core optimization and algorithmic ideas presented are applicable regardless of whether elementwise affines exist, although the performance results may differ.

1. Challenges of Normalization Fusion

Overview: In this section we discuss the challenges of fusing normalization with compute-intensive kernels like GEMMs in the typical way, which fundamentally stem from different tiling strategies. We then present a “naive” fusion solution that forces the GEMM algorithm to obey the same tiling as the normalization algorithm, and observe that it performs well for super small N, but becomes suboptimal or even infeasible as N increases due to tiling constraints and inefficiencies.

Compared with standard activation fusion (e.g. GEMM+ReLU), the fundamental challenge with normalization fusion is the difference in tiling. Normalization is by nature a reduction operation that requires access to data along an entire dimension to compute the correct result. In particular, for LayerNorm and RMSNorm, a typical kernel tiles the input along the outer dimension(s), but not the inner, meaning...