Implementation:NVIDIA TransformerEngine JAX Dense

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

Overview

Provides the dense operation (matrix multiplication with optional bias) as a differentiable JAX function with FP8 quantization and distributed communication support.

Description

The public dense() function delegates to _dense, which uses jax.custom_vjp to define custom forward and backward rules. The forward pass quantizes inputs/kernels via QuantizerSet, applies sharding constraints, and calls tex.gemm with optional collective operations (AllGather/ReduceScatter for tensor parallelism). The backward pass computes gradients for input, kernel, and bias using transposed GEMMs with quantized operands. Supports configurable contracting dimensions for flexible tensor layouts.

The grouped_dense function provides a similar interface for grouped/batched GEMM operations used in Mixture-of-Experts workloads.

This is the core linear transformation building block used by all dense layers in the JAX backend (DenseGeneral, LayerNormDenseGeneral, LayerNormMLP), providing the differentiable FP8 GEMM with integrated distributed communication.

Usage

Use this module when implementing linear transformations in JAX with FP8 quantization support. It is the primary building block for all dense layer operations. Higher-level modules (DenseGeneral, LayerNormDenseGeneral) wrap this function with Flax parameter management.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/jax/dense.py

Lines

1--643

Signature

def dense(
    x: jnp.ndarray,
    kernel: jnp.ndarray,
    bias: jnp.ndarray = None,
    contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((1,), (0,)),
    transpose_batch_sequence: bool = False,
    input_axes: Tuple[str, ...] = None,
    kernel_axes: Tuple[str, ...] = None,
    output_axes: Tuple[str, ...] = None,
    collective_op_set: tex.CollectiveOpSet = tex.noop_collective_op_set,
    quantizer_set: QuantizerSet = noop_quantizer_set,
) -> jnp.ndarray: ...

def grouped_dense(
    x: jnp.ndarray,
    kernel: jnp.ndarray,
    group_sizes: jnp.ndarray,
    contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((1,), (1,)),
    bias: jnp.ndarray = None,
    kernel_amax: jnp.ndarray = None,
    precision: jax.lax.Precision = jax.lax.Precision.DEFAULT,
    preferred_element_type: jnp.dtype = None,
    group_offset: jnp.array = None,
    quantizer_set: QuantizerSet = noop_quantizer_set,
    kernel_fsdp_info: Tuple[str, int] = (None, -1),
) -> jnp.ndarray: ...

Import

from transformer_engine.jax.dense import dense, grouped_dense

I/O Contract

Inputs

Name	Type	Required	Description
x	jnp.ndarray	Yes	Input tensor
kernel	jnp.ndarray	Yes	Weight matrix
bias	jnp.ndarray	No	Optional bias vector
contracting_dims	Tuple[Sequence[int], Sequence[int]]	No	Contracting dimensions for the GEMM (default ((1,), (0,)))
collective_op_set	CollectiveOpSet	No	Collective operations for tensor parallelism
quantizer_set	QuantizerSet	No	FP8 quantizer set for input, kernel, and gradient

Outputs

Name	Type	Description
output	jnp.ndarray	Result of the linear transformation (x @ kernel + bias)

Usage Examples

from transformer_engine.jax.dense import dense
from transformer_engine.jax.quantize import noop_quantizer_set

# Simple dense layer
output = dense(input_tensor, weight_matrix, bias=bias_vector)

# Dense with FP8 quantization and tensor parallelism
output = dense(
    input_tensor, weight_matrix,
    bias=bias_vector,
    quantizer_set=fp8_quantizer_set,
    collective_op_set=tp_collective_ops,
    input_axes=("batch", "hidden"),
    kernel_axes=("hidden", "mlp"),
)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements