Implementation:NVIDIA TransformerEngine Normalization Common Header

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

Overview

Defines the core data structures and class hierarchy for the normalization subsystem, including kernel parameter structs, launch parameters, execution plans, and the normalization plan registry.

Description

normalization/common.h declares the types all normalization operations (LayerNorm, RMSNorm) depend on:

• KernelParamsBase: Base struct holding common fields: rows, cols, pointers to x, mu, rs, gamma, workspace, barrier, and zero_centered_gamma flag.
• ForwardKernelParams: Extends base with z, beta, epsilon, scale/scale_inv/amax pointers, and fp8_out flag.
• BackwardKernelParams: Extends base with dz, dgamma/dbeta partials, and gradient pointers.
• LaunchParams<T>: Wraps kernel params with workspace sizing, barrier bytes, SM count, and stream.
• NormalizationPlanBase: Abstract interface with execute() and getWorkspaceShape().
• TeNormalizationPlan: Concrete plan using custom TE CUDA kernels.
• CudnnNormalizationPlan: Concrete plan using cuDNN frontend graphs.
• NormalizationPlanRegistry: Singleton that caches plans keyed by norm parameters.

Usage

Include this header when implementing normalization backends or extending the normalization system.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/normalization/common.h

Lines

1--395

Signature

namespace transformer_engine { namespace normalization {

struct KernelParamsBase {
  int ctas_per_col, ctas_per_row, rows, cols;
  void *x, *mu, *rs, *gamma, *workspace;
  int *barrier;
  bool zero_centered_gamma;
};

struct ForwardKernelParams : public KernelParamsBase {
  void *z, *beta;
  float epsilon;
  void *scale, *scale_inv, *amax;
  bool fp8_out;
};

struct BackwardKernelParams : public KernelParamsBase {
  void *dz, *dbeta_part, *dgamma_part;
};

template <typename KernelParamsType>
struct LaunchParams {
  size_t workspace_bytes, barrier_bytes, dgamma_part_bytes;
  int multiprocessorCount;
  cudaStream_t stream;
  KernelParamsType params;
};

}}  // namespace

Import

#include "normalization/common.h"

I/O Contract

Inputs

Name	Type	Required	Description
N/A	N/A	N/A	This is a header file defining types and class hierarchy

Outputs

Name	Type	Description
N/A	N/A	Provides type definitions for the normalization subsystem

Usage Examples

#include "normalization/common.h"

using namespace transformer_engine::normalization;

// Set up forward kernel parameters
ForwardKernelParams params;
params.rows = batch_size;
params.cols = hidden_size;
params.epsilon = 1e-5f;
params.fp8_out = true;

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements