Implementation:NVIDIA TransformerEngine JAX XLA Normalization

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, JAX, Normalization
Last Updated	2026-02-07 14:00 GMT

Overview

Implements XLA FFI handlers for LayerNorm and RMSNorm forward and backward passes with optional FP8 quantized output, exposed as custom JAX operations.

Description

GetNormForwardWorkspaceSizes and GetNormBackwardWorkspaceSizes create dummy tensor wrappers and call nvte_layernorm_fwd/nvte_rmsnorm_fwd (or their backward counterparts) to query required workspace sizes. NormForwardFFI extracts input/gamma/beta buffers, constructs TensorWrapper objects with scaling metadata (delayed tensor, MXFP8 block, or no scaling), configures rowwise and optional columnwise (2x2x layout) FP8 output tensors with scale/amax/scale_inv, then dispatches to nvte_layernorm_fwd or nvte_rmsnorm_fwd. The backward handler similarly dispatches to nvte_layernorm_bwd/nvte_rmsnorm_bwd. Both forward and backward handlers have Initialize variants for CUDA graph capture support. SM margin is subtracted from the device's multiprocessor count for normalization kernel configuration.

This extension enables JAX models to use GPU-accelerated fused normalization kernels that combine norm computation with FP8 quantization in a single kernel launch.

Usage

This C++ extension is invoked internally by the Python-side NormFwdPrimitive and NormBwdPrimitive in transformer_engine.jax.cpp_extensions.normalization. Users do not call these FFI handlers directly.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/jax/csrc/extensions/normalization.cpp

Lines

1--383

Signature

namespace transformer_engine {
namespace jax {

pybind11::tuple GetNormForwardWorkspaceSizes(
    size_t batch_size, size_t hidden_size, DType in_dtype,
    DType w_dtype, DType out_dtype, NVTE_Norm_Type norm_type,
    JAXX_Scaling_Mode scaling_mode, bool zero_centered_gamma,
    float epsilon, int sm_margin, bool is_training);

pybind11::tuple GetNormBackwardWorkspaceSizes(
    size_t batch_size, size_t hidden_size, DType in_dtype,
    DType w_dtype, NVTE_Norm_Type norm_type, bool zero_centered_gamma,
    float epsilon, int sm_margin);

Error_Type NormForwardFFI(
    cudaStream_t stream, Buffer_Type input_buf, Buffer_Type gamma_buf,
    Buffer_Type beta_buf, Buffer_Type scale_buf, Buffer_Type amax_buf,
    Result_Type output_buf, Result_Type colwise_output_buf,
    Result_Type mu_buf, Result_Type rsigma_buf,
    Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf,
    Result_Type updated_amax_buf, Result_Type workspace_buf,
    NVTE_Norm_Type norm_type, JAXX_Scaling_Mode scaling_mode,
    JAXX_Quantize_Layout quantize_layout,
    bool zero_centered_gamma, float epsilon, int sm_margin);

Error_Type NormBackwardFFI(...);

} // namespace jax
} // namespace transformer_engine

Import

#include "transformer_engine/normalization.h"
#include <cuda_runtime.h>
#include "../extensions.h"

I/O Contract

Inputs

Name	Type	Required	Description
input_buf	Buffer_Type	Yes	Input tensor buffer
gamma_buf	Buffer_Type	Yes	Scale parameter buffer
beta_buf	Buffer_Type	Yes	Shift parameter buffer (LayerNorm only)
scale_buf	Buffer_Type	Yes	Quantization scale buffer
amax_buf	Buffer_Type	Yes	Amax buffer for delayed scaling
norm_type	NVTE_Norm_Type	Yes	LayerNorm or RMSNorm
scaling_mode	JAXX_Scaling_Mode	Yes	Quantization scaling mode
zero_centered_gamma	bool	Yes	Whether gamma is zero-centered
epsilon	float	Yes	Numerical stability constant

Outputs

Name	Type	Description
output_buf	Result_Type	Normalized output, optionally FP8 quantized
mu_buf	Result_Type	Mean for backward pass (LayerNorm only)
rsigma_buf	Result_Type	Reciprocal std deviation for backward pass
scale_inv_buf	Result_Type	Inverse scale factor for output
updated_amax_buf	Result_Type	Updated amax value

Usage Examples

// This FFI handler is called internally by JAX's XLA compilation pipeline.
// Users interact with it through the Python API:
//   from transformer_engine.jax.cpp_extensions.normalization import normalization_fwd
//   output, mu, rsigma = normalization_fwd(x, gamma, beta, ...)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements