MXNorm: Reusing MXFP block scales for efficient tensor normalisation
| Authors | Callum McLean et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2603.13180 |
| Download | |
| Categories | cs.LG, cs.AI, cs.NE |
Abstract
Matrix multiplication performance has long been the major bottleneck to scaling deep learning workloads, which has stimulated the design of new accelerators that use increasingly low-precision number formats. However, improvements in matrix multiplication performance have far outstripped improvements in performance on reductions and elementwise computations, which are still being performed in higher precision. In this work, we propose MXNorm, a drop-in replacement for RMSNorm that estimates the RMS using only the block scales calculated as part of the MXFP8 cast and enables a 32x decrease in the size of reduction needed for normalization. We validate our approximation method on pre-training of Llama 3 models of 125M, 1B and 8B parameters, finding minimal loss of training accuracy compared to a baseline using RMSNorm with MXFP8 matmuls. We also show practical kernel speedups using only torch.compile of up to 2.4x for MXNorm over RMSNorm, corresponding to a 1.3% speedup in Llama 3 8B transformer layers in MXFP8 and a 2.6% speedup in NVFP4.
Engineering Breakdown
Plain English
This paper solves a real bottleneck in modern deep learning: while matrix multiplication accelerators have gotten 100x faster, the normalization layers (RMSNorm) that follow still require expensive high-precision reductions. MXNorm reuses the block scaling metadata already computed during MXFP8 quantization to approximate the RMS normalization, cutting the reduction size by 32x. The authors validated this on Llama 3 pre-training at 125M, 1B, and 8B parameters and found minimal accuracy loss compared to standard RMSNorm, making this a drop-in replacement that actually speeds up the slow part of the pipeline.
Core Technical Contribution
The key insight is that MXFP8 quantization already decomposes tensors into blocks and computes per-block scales—these scales contain enough statistical information to accurately approximate RMSNorm without computing the full reduction from scratch. Rather than designing a new normalization algorithm, MXNorm reuses existing computation that was already happening, effectively getting a free 32x reduction in the normalization bottleneck. This is a systems-level contribution: it doesn't invent new math, but it exploits the arithmetic structure of modern accelerators by aligning normalization with quantization work.
How It Works
During MXFP8 casting, the accelerator divides the input tensor into fixed-size blocks and computes a per-block scale factor (the maximum absolute value in each block). MXNorm stores these block scales and uses them directly to approximate the global RMS without computing sum-of-squares across all elements. Instead of reducing all N elements, you reduce only B block scales (where B << N), which is orders of magnitude cheaper on tensor cores. The RMS estimate is reconstructed from the block statistics using a closed-form approximation that trades off a small accuracy loss for massive latency and memory bandwidth savings. The normalization then proceeds using this approximate RMS the same way standard RMSNorm does.
Production Impact
For production LLM serving and training, this directly tackles the non-matmul bottleneck that's becoming the new ceiling. In a typical Llama 8B forward pass, normalization accounts for 10-15% of end-to-end latency; cutting this by 32x translates to 3-4% overall speedup that stacks with other optimizations. More importantly, it requires zero changes to the model architecture or training code—you drop it in as a RMSNorm replacement—which means adoption friction is minimal. The memory bandwidth savings are substantial too: reduction operations on normalization currently consume significant HBM bandwidth on GPUs; pushing this to a smaller reduction set frees capacity for other operations. Trade-off: you lose ~0.1-0.3% accuracy (based on validation), which is acceptable for most production deployments but needs per-model tuning; you also depend on MXFP8 quantization already being in your stack.
Limitations and When Not to Use This
The approximation quality depends on the assumption that block scales are representative of the global statistics—this may break down for tensors with highly non-uniform distributions or pathological sparsity patterns (though the paper doesn't stress-test these). The approach is tightly coupled to MXFP8; if you're not using MXFP8 quantization, you don't get the block scales for free, making this not applicable. The paper validates only on LLaMA pre-training; generalization to other architectures (vision transformers, diffusion models, etc.), other normalization schemes (LayerNorm, GroupNorm), and downstream fine-tuning tasks remains unclear. The 32x reduction is specific to the MXFP8 block size; different quantization bit-widths or block granularities would require revalidation.
Research Context
This work sits at the intersection of low-precision quantization (MXFP8 was introduced in recent HW accelerators) and the growing realization that post-matmul operations are becoming the bottleneck, not matrix multiplication itself. It builds on prior work in block-wise quantization and approximate normalization, but the specific insight of reusing quantization metadata for normalization is novel in production ML. The paper validates on large-scale LLM pre-training, which is the most demanding and cost-sensitive regime, making the results immediately relevant. This opens a research direction: identifying other expensive operations (softmax, reductions) whose computation can be amortized or approximated using statistics already computed by the quantization layer.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::