Implementation:NVIDIA TransformerEngine JAX Cpp GEMM
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements JAX custom primitives for general matrix multiplication (GEMM) and grouped GEMM operations, with FP8 quantization support and collective communication overlapping for distributed training.
Description
GemmPrimitive wraps cuBLAS GEMM calls via FFI, handling FP8/FP4 scaled tensor inputs with proper scale factor management. GroupedGemmPrimitive supports batched GEMM for Mixture-of-Experts workloads. CollectiveOp and CollectiveOpSet define AllGather/ReduceScatter operations that can be overlapped with GEMM computation for tensor-parallel training. The module handles operand quantization, layout transposition, and scale swizzling for block-scaled tensors. It uses JAX's custom_partitioning for SPMD-aware sharding.
This is one of the most critical performance components -- all linear transformations in the transformer flow through these GEMM primitives, making this the primary interface to cuBLAS for FP8/FP4 accelerated matrix multiplication.
Usage
Use this module indirectly through dense(), layernorm_dense(), or layernorm_mlp(). Direct usage is needed when implementing custom linear layers with FP8 quantization or when configuring collective communication overlap for tensor parallelism.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/cpp_extensions/gemm.py- Lines
- 1--2094
Signature
class CollectiveOp(Enum):
NONE = 0
ALL_GATHER = 1
REDUCE_SCATTER = 2
class CollectiveOpSet:
input_op: CollectiveOp = CollectiveOp.NONE
output_op: CollectiveOp = CollectiveOp.NONE
mesh_axis: str = ""
class GemmPrimitive(BasePrimitive): ...
class GroupedGemmPrimitive(BasePrimitive): ...
def gemm(
lhs: Union[jnp.ndarray, ScaledTensor],
rhs: Union[jnp.ndarray, ScaledTensor],
dimension_numbers: Tuple,
bias: Optional[jnp.ndarray] = None,
fuse_bias: bool = False,
...,
) -> jnp.ndarray: ...
def grouped_gemm(
lhs: Union[jnp.ndarray, ScaledTensor],
rhs: Union[jnp.ndarray, ScaledTensor],
group_sizes: jnp.ndarray,
...,
) -> jnp.ndarray: ...
Import
from transformer_engine.jax.cpp_extensions.gemm import gemm, grouped_gemm, CollectiveOpSet
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| lhs | Union[jnp.ndarray, ScaledTensor] |
Yes | Left-hand side operand, optionally FP8 quantized |
| rhs | Union[jnp.ndarray, ScaledTensor] |
Yes | Right-hand side operand, optionally FP8 quantized |
| dimension_numbers | Tuple |
Yes | Contracting and batch dimension specification |
| bias | Optional[jnp.ndarray] |
No | Optional bias to fuse with GEMM |
| collective_op_set | CollectiveOpSet |
No | Optional collective operations for tensor parallelism |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Matrix multiplication result |
Usage Examples
from transformer_engine.jax.cpp_extensions.gemm import gemm, CollectiveOpSet, CollectiveOp
# Standard GEMM
output = gemm(input_tensor, weight, dimension_numbers=((1,), (0,)))
# GEMM with tensor-parallel communication overlap
collective = CollectiveOpSet(
input_op=CollectiveOp.ALL_GATHER,
output_op=CollectiveOp.NONE,
mesh_axis="tp"
)
output = gemm(input_tensor, weight, dimension_numbers=((1,), (0,)),
collective_op_set=collective)