Implementation:NVIDIA TransformerEngine Float8 Storage

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

Overview

Low-level storage class for Float8 (FP8) tensor data, providing dequantization, transpose caching, and usage-based data management.

Description

Float8TensorStorage is a mixin class that holds the data attributes of Float8Tensor: FP8 data, scale inverse, FP8 dtype, transpose cache, and optional quantizer. It provides:

• Dequantization via _FromFloat8Func (custom autograd function using tex.dequantize)
• Transpose caching with lazy creation and invalidation tracking
• Usage management (update_usage) that creates or removes rowwise/columnwise data based on GEMM requirements
• Serialization helpers (prepare_for_saving/restore_from_saved) for autograd compatibility

When instantiated directly (not through Float8Tensor), it has lower CPU overhead but less functionality. It supports non-TN FP8 GEMM detection for deciding when transpose data is needed.

Usage

Used internally by the Float8 tensor system. Direct instantiation is for performance-critical internal usage only.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py

Lines

1--236

Signature

class Float8TensorStorage(QuantizedTensorStorage):
    def __new__(cls, *args, data, fp8_scale_inv, fp8_dtype, data_transpose=None, quantizer=None, **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) -> Float8TensorStorage: ...

Import

from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage

I/O Contract

Inputs

Name	Type	Required	Description
data	Optional[torch.Tensor]	Yes	FP8 data tensor
fp8_scale_inv	torch.Tensor	Yes	Inverse scaling factor
fp8_dtype	TE_DType	Yes	FP8 data type (e.g., E4M3, E5M2)
data_transpose	Optional[torch.Tensor]	No	Cached FP8 transpose
quantizer	Optional[Quantizer]	No	Quantizer used to create this tensor

Outputs

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

Usage Examples

from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage

storage = Float8TensorStorage(
    data=fp8_data,
    fp8_scale_inv=scale_inv,
    fp8_dtype=TE_DType.kFloat8E4M3,
)
hp_tensor = storage.dequantize(dtype=torch.bfloat16)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements