Implementation:NVIDIA TransformerEngine JAX Hadamard
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements Randomized Hadamard Transform (RHT) utilities used in NVFP4 weight gradient quantization to improve quantization quality.
Description
get_rht_matrix constructs a 16x16 Hadamard matrix using scipy.linalg.hadamard, multiplies it by a fixed random sign vector (from get_wgrad_sign_vector), and normalizes by 1/sqrt(16). apply_rht reshapes the input into blocks of 16 and multiplies by this matrix. get_sign_from_vector converts the sign vector to a bitmask integer for use in C++ kernels. The inverse RHT uses the matrix inverse for dequantization.
This module enables higher-quality NVFP4 quantization by applying a randomized orthogonal transformation before quantization, which distributes outlier values more evenly across the quantization blocks.
Usage
Use this module when working with NVFP4 quantization that requires the Randomized Hadamard Transform. It is called internally by NVFP4Quantizer and NVFP4Dequantizer.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/quantize/hadamard.py- Lines
- 1--46
Signature
def get_wgrad_sign_vector() -> list[int]:
"""Get a fixed sign vector for the RHT used in NVFP4 weight gradient quantization."""
...
def get_sign_from_vector(vector: list[int]) -> int:
"""Convert a sign vector to a bitmask integer."""
...
def apply_rht(x: jnp.ndarray, inverse=False) -> jnp.ndarray:
"""Apply the Randomized Hadamard Transform (RHT) to the input tensor."""
...
def get_rht_matrix() -> jnp.ndarray:
"""Get the 16x16 RHT matrix pre-multiplied by the random sign mask."""
...
Import
from transformer_engine.jax.quantize.hadamard import apply_rht, get_rht_matrix, get_sign_from_vector
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| x | jnp.ndarray |
Yes | Input tensor (last dimension must be divisible by 16) |
| inverse | bool |
No | Whether to apply the inverse transform (default False) |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Transformed tensor with same shape as input |
Usage Examples
from transformer_engine.jax.quantize.hadamard import apply_rht, get_sign_from_vector, get_wgrad_sign_vector
import jax.numpy as jnp
# Apply RHT before NVFP4 quantization
transformed = apply_rht(weight_gradient)
# Apply inverse RHT after dequantization
restored = apply_rht(dequantized_data, inverse=True)
# Get sign bitmask for C++ kernel
sign_mask = get_sign_from_vector(get_wgrad_sign_vector())