Implementation:NVIDIA TransformerEngine LayerNorm API

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

Overview

Implements the C++ API entry points for LayerNorm forward and backward passes, including input validation, backend selection between TE and cuDNN, workspace allocation, and optional FP8 transpose.

Description

normalization/layernorm/ln_api.cpp bridges the public C API (nvte_layernorm_fwd/bwd) to the internal plan-based execution system:

• layernorm_fwd: Validates tensor shapes (x must be 2D, gamma/beta same shape, mu/rsigma correct shapes), selects between TE and cuDNN backends (preferring cuDNN for MXFP8 scaling), obtains a cached normalization plan from NormalizationPlanRegistry, and executes it. On first call with empty workspace, returns required workspace shape. Optionally computes FP8 transpose for columnwise data.
• layernorm_bwd: Validates backward inputs, selects backend, dispatches backward computation for dgamma, dbeta, and dx.
• layernorm_bwd_add: Fused backward variant that adds gradient to dx.

Supported type combinations: FP32, FP16, BF16 inputs with FP32 compute, and optional FP8/MXFP8 output.

Usage

This is the primary LayerNorm entry point, called by the C API wrapper functions which in turn are called by framework bindings.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/normalization/layernorm/ln_api.cpp

Lines

1--230

Signature

namespace transformer_engine {

void layernorm_fwd(const Tensor& x, const Tensor& gamma, const Tensor& beta,
                   const float epsilon, Tensor* z, Tensor* mu, Tensor* rsigma,
                   Tensor* workspace, const int multiprocessorCount,
                   const bool zero_centered_gamma, cudaStream_t stream);

void layernorm_bwd(const Tensor& dz, const Tensor& x, const Tensor& mu,
                   const Tensor& rsigma, const Tensor& gamma,
                   Tensor* dx, Tensor* dgamma, Tensor* dbeta,
                   Tensor* workspace, const int multiprocessorCount,
                   const bool zero_centered_gamma, cudaStream_t stream);

}  // namespace transformer_engine

Import

#include <transformer_engine/normalization.h>

I/O Contract

Inputs

Name	Type	Required	Description
x	Tensor&	Yes	Input tensor [N, H]
gamma	Tensor&	Yes	Gamma weight [H]
beta	Tensor&	Yes	Beta weight [H]
epsilon	float	Yes	Numerical stability constant (must be >= 0)
multiprocessorCount	int	Yes	Number of GPU SMs

Outputs

Name	Type	Description
z	Tensor*	Normalized output [N, H] (optionally FP8)
mu	Tensor*	Per-row mean [N]
rsigma	Tensor*	Per-row inverse std deviation [N]

Usage Examples

#include <transformer_engine/normalization.h>

// Query workspace size (first call with empty workspace)
Tensor workspace;
layernorm_fwd(x, gamma, beta, 1e-5f, &z, &mu, &rsigma,
              &workspace, sm_count, false, stream);
// workspace.data.shape is now set to required size

// Allocate workspace and execute
// ... allocate workspace.data.dptr ...
layernorm_fwd(x, gamma, beta, 1e-5f, &z, &mu, &rsigma,
              &workspace, sm_count, false, stream);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements