Implementation:NVIDIA TransformerEngine JAX LayerNorm Dense
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Normalization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements a fused layer normalization followed by dense (GEMM) operation as a single differentiable JAX function with FP8 quantization and sharding support.
Description
The public layernorm_dense() delegates to _layernorm_dense via jax.custom_vjp. The forward pass performs normalization (tex.normalization_fwd), applies sharding constraints, quantizes the normalized output and kernel, then executes GEMM (tex.gemm) with optional bias addition. The backward pass computes gradients for all inputs (x, kernel, gamma, beta, bias) using the saved forward pass context, with proper quantization of intermediate gradients for FP8 weight gradient computation.
This is a key fusion pattern for transformer architectures -- combining normalization and linear projection into a single operation reduces memory traffic by avoiding materialization of the normalized intermediate tensor. It is used by LayerNormDenseGeneral in the Flax modules.
Usage
Use this function when implementing fused layer normalization + linear projection layers. It is the primary building block for LayerNormDenseGeneral in Flax modules.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/layernorm_dense.py- Lines
- 1--354
Signature
def layernorm_dense(
x: jnp.ndarray,
kernel: jnp.ndarray,
gamma: jnp.ndarray,
beta: jnp.ndarray,
bias: jnp.ndarray = None,
norm_type: str = "layernorm",
zero_centered_gamma: bool = False,
epsilon: float = 1e-6,
transpose_batch_sequence: bool = False,
layernorm_input_axes: Tuple[str, ...] = None,
dot_input_axes: Tuple[str, ...] = None,
kernel_axes: Tuple[str, ...] = None,
quantizer_set: QuantizerSet = noop_quantizer_set,
) -> jnp.ndarray: ...
Import
from transformer_engine.jax.layernorm_dense import layernorm_dense
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| x | jnp.ndarray |
Yes | Input tensor |
| kernel | jnp.ndarray |
Yes | Weight matrix for the linear transformation |
| gamma | jnp.ndarray |
Yes | Layer normalization scale parameter |
| beta | jnp.ndarray |
Yes | Layer normalization shift parameter |
| bias | jnp.ndarray |
No | Optional bias for the linear transformation |
| norm_type | str |
No | Normalization type: "layernorm" or "rmsnorm" (default "layernorm") |
| zero_centered_gamma | bool |
No | Whether gamma is zero-centered (default False) |
| epsilon | float |
No | Numerical stability constant (default 1e-6) |
| quantizer_set | QuantizerSet |
No | FP8 quantizer set |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Result of LayerNorm(x) @ kernel + bias |
Usage Examples
from transformer_engine.jax.layernorm_dense import layernorm_dense
# Fused LayerNorm + Dense forward pass
output = layernorm_dense(
x=input_tensor,
kernel=weight_matrix,
gamma=ln_scale,
beta=ln_bias,
bias=dense_bias,
norm_type="layernorm",
epsilon=1e-5,
quantizer_set=fp8_quantizer_set,
)