Implementation:NVIDIA TransformerEngine Ops L2Normalization

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

Overview

Fusible operation that applies L2 normalization over the last dimension, scaling each vector to unit L2 norm with optional JIT warmup.

Description

L2Normalization is a parameter-free BasicOperation that computes y = x / sqrt(sum(x^2) + eps) using JIT-fused implementations. The forward pass computes and caches rsqrt_norm for efficient backward computation. It supports JIT warmup for common hidden sizes used in query-key normalization (32, 64, 80, 96, 128, 256) and CPU activation offloading. The implementation uses torch.compile-compatible fused functions.

Usage

Used for query-key normalization in attention mechanisms, particularly for QK-norm variants that require unit-norm queries and keys.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/basic/l2normalization.py

Lines

1--130

Signature

class L2Normalization(BasicOperation):
    def __init__(self, *, eps: float = 1e-6, seq_length: Optional[int] = None, micro_batch_size: Optional[int] = None) -> None: ...
    def op_forward(self, ctx, input_, prev_op_grad_output_quantizer, next_op_input_quantizer) -> torch.Tensor: ...
    def op_backward(self, ctx, grad_output) -> Tuple[torch.Tensor, Tuple[()]]: ...

Import

from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization

I/O Contract

Inputs

Name	Type	Required	Description
input_	torch.Tensor	Yes	Input tensor to normalize
eps	float	No	Numerical stability constant (default 1e-6)

Outputs

Name	Type	Description
output	torch.Tensor	L2-normalized tensor (unit norm along last dimension)

Usage Examples

from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization

l2norm = L2Normalization(eps=1e-6)
normalized = l2norm.op_forward(ctx, query_tensor, None, None)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements