MXNorm: Reusing MXFP block scales for efficient tensor normalisation

Callum McLean; Luke Y. Prince; Alexandre Payot; Paul Balança; Carlo Luschi

cs.LGMar 13, 2026

MXNorm: Reusing MXFP block scales for efficient tensor normalisation

Authors: Callum McLean, Luke Y. Prince, Alexandre Payot, Paul Balança, Carlo Luschi

Organizations: Graphcore

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.

Figures & tables

GPU	dtype	Tensor Core		CUDA Core		Mem. BW
		(TFLOPS)		(TFLOPS)		(TB/s)
V100	FP16	$125$		$15.7$		$0.9$
A100	FP16	$312$	[2.5x]	$19.5$	[1.2x]	$2.0$	[2.2x]
H100	FP8	$1979$	[15.8x]	$67.0$	[4.3x]	$3.4$	[3.7x]
GB200	FP4	$10\,000$	[80.0x]	$80.0$	[5.1x]	$8.0$	[8.9x]
Rubin*	FP4	$35\,000$	[280x]	$130.0$	[8.3x]	$22.0$	[24.4x]

Table 1: Comparison of GPU performance and memory bandwidth. Improvement relative to V100 shown in square brackets. dtype: Lowest supported precision of floating point format for matrix multiplication. *Anticipated specification of Rubin generation GPUs that have yet to be released commercially.

Figure 1: Computational graphs for RMSNorm, MXCast, and MXNorm in the context of Norm + Linear layer pattern. Top left: RMSNorm + Linear graph for high precision training. Top middle: RMSNorm + Linear graph with linear inputs cast to MX (MXLinear). Top right: RMSNorm approximated with MXNorm and RMSNorm weight fused with Linear weight. Bottom left: RMSNorm graph. Bottom middle: MXCast graph for MXFP8 activations. Note that MXCast is applied to weights (E4M3 values) and gradients (E4M3 or E5M2 values) as well. Bottom right: MXNorm graph.

Side by side diff-like comparison of RMSNorm + MXCast against MXNorm without norm gain parameters. Input

X

with shape

[T\times K\times B]

is received in blocked format where it originally had

D

columns that have been partitioned in

K

blocks of size

B

. We omit the use of a normalization stability constant

\epsilon

for brevity. Block absmaxes are divided by 256 (the largest representable power of 2 in E4M3) to ensure maximal use of range.

p

is a scalar norm exponent, used with p=1 for the arithmetic mean and p=2 for the RMS (quadratic mean) of absmaxes.

c

is a scalar correction factor dependent on block size

B

and

p

to ensure

\tilde{\rho}\approx\rho

when

X_{t}

are i.i.d. samples from a Gaussian distribution.

\text{cast}(\cdot;EaMb)

represents a cast to floating point format

\mathbb{F}_{EaMb}

with

a

exponent and

b

mantissa bits. E8M0 casting uses an implementation defined power-of-2 rounding function

2^{\lfloor\log_{2}\cdot\rceil}

[ 19 ] .

\mathbb{F}_{\star}

corresponds to a higher-precision floating point format.

Figure 2: MXNorm as an approximation of RMSNorm. Left: MX scale distribution of normalized tensors. Middle: MX value distribution of normalized tensors. Right: MXNorm

r^{2}

goodness-of-fit approaches 1 with more blocks.

Figure 3: Learning rate sensitivity of MXNorm compared to RMSNorm. Left: 125M parameter model (depth=8, width=1024). Right: 1B parameter model (depth=16, width=2048).

Figure 4: Training loss convergence of 8B parameter models trained on 300B tokens with MXNorm and RMSNorm

Norm Type	ARC-C	ARC-E	BoolQ	CSQA	HellaSwag	MMLU	OBQA	PIQA	SIQA	WinoGr
RMSNorm	45.3	75.9	73.2	68.5	70.4	37.9	50.4	74.8	56.5	66.8
MXNorm (p=1)	39.1	73.7	56.6	62.4	68.1	37.2	49.8	74.5	52.5	64.4
MXNorm (p=2)	45.8	74.1	72.6	67.1	71.5	38.0	54.2	76.2	55.5	65.2

Table 2: Zero-shot performance on OLMES natural language tasks for 8B parameter models

Figure 5: Example outlier feature that appears at the same step as the first loss spike in a 8B parameter training run with MXNorm using mean over block absmax to estimate RMS.

Figure 6: RMS of MXNorm outputs tracked throughout training across all 16 layers of small-scale proxy used to reproduce loss spikes. Envelope of loss spike is visible as a deflection below 1.0 just before

10^{9}

tokens. Prior to the loss spike, normalization by MXNorm closely matches behaviour of RMSNorm.

Figure 7: Histograms of input and output tensor elements for RMSNorm and MXNorm using mean or RMS over block absmaxes to compute normalization scale.

Figure 8: Relative speedup of isolated MXNorm over RMSNorm-MXCast. Normalization and casting have been fused together into a single kernel in both cases.

B

denotes the block size.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A.1 : Representation of a block of 32 elements stored in MXFP8.

\text{E}x\text{M}y

denotes the

x

bits allocated to E xponent and

y

bits allocated to M antissa. Each contiguous block of 32 elements in an MXFP8 tensor is represented by a tuple of 32 E4M3 values along with an unsigned E8M0 scale shared among the 32 elements. During dequantization, each value is multiplied by its corresponding shared scale.

	Model Size
	125M	1B	8B
Global batch size	32	256	256
Sequence length	4096	4096	4096
Total training tokens	2.62B	20.97B	314.57B
Transformer block count	8	16	32
Hidden dimension	1024	2048	4096

Appendix

Table 3: Pretraining configuration for Llama 3 models

	Set of tested values
Tokens	$\{2^{12},2^{13},2^{14},2^{15},2^{16}\}$
Hidden dim	$\{k\cdot 2^{r}:k\in\{1,3,5,7\},r\in\mathbb{N}\}\cap[2^{10},2^{14}]$
Block size	$\{16,32,64\}$
MXFP dtype	{E4M3, E2M1}

Appendix

Table 4: Kernel benchmark cases. All possible combinations are profiled.

	Set of tested values
Batch size	$\{4,8\}$
Tokens	$\{2^{12},2^{13},2^{14},2^{15},2^{16}\}$
MXFP type	$\{\text{MXFP8},\text{NVFP4}\}$

Appendix

Table 5: Layer benchmark cases

Figure A.2 : Training loss convergence of 1B parameter models trained on 20B tokens with MXNorm and RMSNorm.

Figure A.3 : Boxplot of spike scores showing the mean, interquartile range, and range of values over 8 training runs. Spike score is defined as the area under a detrended loss curve above 3 standard deviations.

Figure A.4 : The post-round MXNorm scheme training loss shows instability that significantly diverges from the RMSNorm baseline.

Layer Normalization

Layer normalization centers and scales neural activations to stabilize training, shape representations, and condition inputs across layers.

See all 29 papers

SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization
Aug 10, 2026Gyudong Kim, Wonjun Han, Young Geun KimGPU Kernel Optimization Tensor Parallelism
Department of Computer Science and Engineering, Korea University, Seoul 02855, Republic of Korea
RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization
Mar 20, 2026Shenyang Deng, Zhuoli Ouyang, Tianyu Pang +4Neural Network Optimization Muon Optimizer
Dartmouth College · International Computer Science Institute · University of California, Berkeley +1
Format-Aware Fusion for Fast FP4 Pretraining
Sep 4, 2026Robert HuLow-Precision Training LLM Quantization
GRAPHCORE RESEARCH

Aug 10, 2026J/K move · Enter open · S save

SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization

Gyudong Kim, Wonjun Han, Young Geun Kim

Department of Computer Science and Engineering, Korea University, Seoul 02855, Republic of Korea

GPU Kernel Optimization Tensor Parallelism LLM Inference Acceleration Distributed Training

Open paper AskPDF arXiv

Query-Key Normalization (QK-Norm) improves the training stability and quality of modern Large Language Models (LLMs). However, under Tensor Parallelism (TP), layerwise QK-Norm introduces additional cross-GPU communication because the normalization factor depends on the full hidden vector. We present SwiftQK, a multi-GPU RMSNorm kernel that exchanges only scalar normalization statistics and overlaps the remaining Peer-to-Peer reduction with independent element-wise computation in a deadlock-safe persistent kernel. Evaluations on recent LLMs show that SwiftQK reduces QK-Norm latency by 81.4--93.9% relative to the standard TP QK-Norm using full-vector All-Gather. In end-to-end serving, SwiftQK reduces TPOT on average by 29.5% over the All-Gather-based baseline and by 14.3% over an optimized scalar-aggregation implementation.

Abstract

Layer Normalization

Explore similar work

SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization

RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization

Format-Aware Fusion for Fast FP4 Pretraining