Implementation:NVIDIA TransformerEngine MXFP8 Storage

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

Overview

Low-level storage class for MXFP8 (Microscaling FP8) tensor data, supporting separate rowwise and columnwise scaled data with GEMM-swizzled scales.

Description

MXFP8TensorStorage is a mixin class for MXFP8Tensor that holds separate rowwise and columnwise FP8 data with independent scaling factors. Unlike Float8TensorStorage which uses a single global scale, MXFP8 uses block-wise scaling (one scale per block of elements). Key features:

• Dual data paths -- Separate rowwise and columnwise scaled data and scale inverses
• GEMM swizzled scales -- Flag indicating whether scales are in the format expected by GEMM kernels
• Dequantization via _FromMXFP8Func (custom autograd function)
• View operations that validate the inner dimension is unchanged (MXFP8 cannot reshape the inner dimension)
• Usage management that only disables data (cannot generate columnwise from rowwise)

Usage

Used internally by the MXFP8 tensor system. Columnwise data can only be created by x2 scaling kernels during quantization.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py

Lines

1--277

Signature

class MXFP8TensorStorage(QuantizedTensorStorage):
    def __new__(cls, rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer, with_gemm_swizzled_scales, *args, **kwargs): ...
    def dequantize(self, *, dtype=torch.float32) -> torch.Tensor: ...
    def update_usage(self, rowwise_usage=None, columnwise_usage=None): ...
    def get_data_tensors(self, rowwise_data=True, columnwise_data=True): ...
    def clear(self): ...
    def view(self, shape) -> MXFP8TensorStorage: ...

Import

from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage

I/O Contract

Inputs

Name	Type	Required	Description
rowwise_data	Optional[torch.Tensor]	Yes	Row-scaled FP8 data
rowwise_scale_inv	Optional[torch.Tensor]	Yes	Row-wise scaling factors
columnwise_data	Optional[torch.Tensor]	No	Column-scaled FP8 data
columnwise_scale_inv	Optional[torch.Tensor]	No	Column-wise scaling factors
fp8_dtype	TE_DType	Yes	FP8 data type
quantizer	Optional[Quantizer]	No	Quantizer used to create this tensor
with_gemm_swizzled_scales	bool	Yes	Whether scales are GEMM-ready

Outputs

Name	Type	Description
dequantized	torch.Tensor	High-precision tensor from dequantize()

Usage Examples

from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage

storage = MXFP8TensorStorage(
    rowwise_data=mxfp8_data,
    rowwise_scale_inv=row_scales,
    columnwise_data=None,
    columnwise_scale_inv=None,
    fp8_dtype=TE_DType.kFloat8E4M3,
    quantizer=mxfp8_quantizer,
    with_gemm_swizzled_scales=True,
)
hp_tensor = storage.dequantize(dtype=torch.bfloat16)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements