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
10000
[80.0x]
80.0
[5.1x]
8.0
[8.9x]
Rubin*
FP4
35000
[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×K×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 ϵ 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 ρ~≈ρ when Xt are i.i.d. samples from a Gaussian distribution. cast(⋅;EaMb) represents a cast to floating point format FEaMb with a exponent and b mantissa bits. E8M0 casting uses an implementation defined power-of-2 rounding function 2⌊log2⋅⌉ [ 19 ] . F⋆ 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 r2 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 109 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. ExMy 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
{212,213,214,215,216}
Hidden dim
{k⋅2r:k∈{1,3,5,7},r∈N}∩[210,214]
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
{212,213,214,215,216}
MXFP type
{MXFP8,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.
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.
Gyudong Kim, Wonjun Han, Young Geun Kim
Department of Computer Science and Engineering, Korea University, Seoul 02855, Republic of Korea
Preconditioned adaptive methods have gained significant attention for training deep neural networks, as they capture rich curvature information of the loss landscape. The central challenge in this field lies in balancing preconditioning effectiveness with computational efficiency of implementing the preconditioner. Among recent advances, Muon stands out by using Newton-Schulz iteration to obtain preconditioned updates without explicitly constructing the preconditioning matrix. Despite its advantages, the efficiency of Muon still leaves room for further improvement. In this paper, we introduce RMNP (Row Momentum Normalized Preconditioning), an optimizer that replaces Newton-Schulz iteration with a simple row-wise (din) ℓ2 normalization operation, motivated by the empirically observed diagonal block structure of the Transformer layerwise Hessian. We empirically verified that orthogonalization and row-wise (on input dim) ℓ2 normalization are asymptotically equivalent in the case of the transformer. This substitution reduces the per-iteration computational complexity from O(mn⋅min(m,n)) to O(mn) for an m×n weight matrix while maintaining comparable optimization performance. Theoretically, we establish convergence guarantees for RMNP in the non-convex setting that match recent results for Muon optimizers, achieving the minimax optimal complexity. Extensive experiments on large language model pretraining show that RMNP delivers competitive optimization performance compared with Muon while substantially reducing preconditioning wall-clock time. Our code is available at https://github.com/Dominator-Index/RMNP.
Shenyang Deng, Zhuoli Ouyang, Tianyu Pang +4
Dartmouth College · International Computer Science Institute · University of California, Berkeley +1
Four-bit floating-point (FP4) Tensor Cores accelerate matrix multiplication, but scale computation, operand packing, layout construction, and saved backward state can erase the gain. We present format-aware fusion, which co-designs each quantization producer with its scale domain and consumer layout for native \mxfp{}, global \nvfp{}, and cooperative-thread-array-local \nvfp{}. We evaluate Llama-3-family 8B pretraining through 160 billion tokens using bfloat16 output projections and compiled cross entropy. In matched same-accelerator probes, bfloat16 and Transformer Engine \nvfp{} reach 18.8K and 27.6K tokens/s/GPU, while our fastest custom route reaches 37.9K. \mxfp{} with row-gradient stochastic rounding and fixed-sign 32-value Hadamard weight-gradient preconditioning reaches 37.2K tokens/s/GPU (86.3% bfloat16 model FLOP utilization) and ends 2.11% above the raw bfloat16 training-loss endpoint. A Transformer Engine recipe with four final bfloat16 blocks ends 0.87% above bfloat16 at 27.1K tokens/s/GPU. Downstream rankings differ from training-loss rankings, showing that FP4 outcomes depend jointly on scale contract, operand, and execution path.