Implementation:NVIDIA TransformerEngine PyTorch Ext Normalization
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements LayerNorm and RMSNorm forward and backward passes with optional fused output quantization to FP8/MXFP8/NVFP4 formats.
Description
Forward functions (layernorm_fwd, rmsnorm_fwd) select between four implementation paths based on the quantizer type: UNFUSED (norm then separate quantize), FULLY_FUSED (single kernel for delayed scaling/MXFP8), FUSED_NORM_AMAX_FP8 (compute norm + amax then quantize for current scaling), and FUSED_NORM_AMAX_NVFP4 (similar for NVFP4). Uses two-pass execution (query workspace size, allocate, execute) for the NVTE kernels. Backward functions (layernorm_bwd, rmsnorm_bwd) compute dx, dgamma, and dbeta gradients. Supports zero-centered gamma and SM margin control for occupancy tuning. The cuDNN MXFP8 path requires 128x128 tile alignment.
Usage
Normalization layers are present in every Transformer block. Fusing normalization with FP8 quantization eliminates a separate quantization pass and reduces memory traffic significantly.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/csrc/extensions/normalization.cpp- Lines
- 1--449
Signature
namespace transformer_engine::pytorch {
py::object layernorm_fwd(at::Tensor input, at::Tensor weight, at::Tensor bias,
float eps, py::handle quantizer, int sm_margin, bool zero_centered_gamma);
py::object rmsnorm_fwd(at::Tensor input, at::Tensor weight,
float eps, py::handle quantizer, int sm_margin, bool zero_centered_gamma);
std::vector<at::Tensor> layernorm_bwd(at::Tensor dz, at::Tensor x,
at::Tensor mu, at::Tensor rsigma, at::Tensor gamma,
int sm_margin, bool zero_centered_gamma);
std::vector<at::Tensor> rmsnorm_bwd(at::Tensor dz, at::Tensor x,
at::Tensor rsigma, at::Tensor gamma,
int sm_margin, bool zero_centered_gamma);
}
Import
#include "../extensions.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| input | at::Tensor |
Yes | Input tensor to normalize |
| weight | at::Tensor |
Yes | Normalization weight (gamma) |
| bias | at::Tensor |
No | Normalization bias (beta), LayerNorm only |
| eps | float |
Yes | Epsilon for numerical stability |
| quantizer | py::handle |
No | Optional quantizer for fused FP8 output |
| zero_centered_gamma | bool |
No | Whether gamma is zero-centered |
Outputs
| Name | Type | Description |
|---|---|---|
| output | py::object |
Normalized (and optionally quantized) tensor |
| mu | at::Tensor |
Mean (LayerNorm only), saved for backward |
| rsigma | at::Tensor |
Reciprocal standard deviation, saved for backward |
Usage Examples
import transformer_engine_torch as tex
# Fused LayerNorm + FP8 quantization
output, mu, rsigma = tex.layernorm_fwd(
input, weight, bias, eps=1e-5,
quantizer=fp8_quantizer,
sm_margin=0, zero_centered_gamma=False
)