`core.inference.quantization.mxfp8_tensor`#

Module Contents#

Classes#

MXFP8Tensor

MXFP8 tensor wrapper storing quantized data and E8M0 scale bytes.

Functions#

`_ceil_div`
`ensure_mxfp8_scale_dtype`	Return the E8M0 view required by PyTorch scaled GEMM APIs.
`validate_mxfp8_tensor`	Thorough MXFP8 validation for GEMM consumers, courtesy of Codex.

Data#

`MXFP8Backend`
`_MXFP8_SCALE_DTYPES`

API#

core.inference.quantization.mxfp8_tensor._ceil_div(a, b)#

core.inference.quantization.mxfp8_tensor.ensure_mxfp8_scale_dtype(scale: torch.Tensor) → torch.Tensor#

Return the E8M0 view required by PyTorch scaled GEMM APIs.

FlashInfer exposes the same MXFP8 scale bytes as uint8.

core.inference.quantization.mxfp8_tensor.MXFP8Backend#: None

core.inference.quantization.mxfp8_tensor._MXFP8_SCALE_DTYPES: dict[core.inference.quantization.mxfp8_tensor.MXFP8Backend, torch.dtype]#: None

core.inference.quantization.mxfp8_tensor.validate_mxfp8_tensor( tensor: MXFP8Tensor, *, expected_backend: Optional[core.inference.quantization.mxfp8_tensor.MXFP8Backend] = None, tensor_name: str = 'MXFP8 tensor', ) → None#: Thorough MXFP8 validation for GEMM consumers, courtesy of Codex.

class core.inference.quantization.mxfp8_tensor.MXFP8Tensor#

MXFP8 tensor wrapper storing quantized data and E8M0 scale bytes.

data: torch.Tensor#: None

scale: torch.Tensor#: None

backend: Optional[core.inference.quantization.mxfp8_tensor.MXFP8Backend]#: None

dtype: Optional[torch.dtype]#: ‘field(…)’

property shape: torch.Size#: Shape of the quantized data storage.

property device: torch.device#: Device holding the quantized tensor.

size(idx: Optional[int] = None)#: Wrapper for calling self.data.size()

scale_2d(K: Optional[int] = None) → torch.Tensor#

Reshape 1D swizzled scale to 2D for scaled_grouped_mm / scaled_mm.

Swizzle pads rows to multiples of 128 and cols to multiples of 4. Returns (padded_M, padded_cols) where padded_cols = ceil(K//32, 4) * 4.

quantize_( value: torch.Tensor, ) → core.inference.quantization.mxfp8_tensor.MXFP8Tensor#

Quantize a logical tensor into existing storage without changing pointers.

The source dtype is preserved when supported by the configured backend so the quantizer does not introduce an avoidable intermediate downcast. FlashInfer FP32 inputs are converted to BF16 because that backend accepts FP16/BF16 inputs only. For legacy instances with unknown logical dtype, the first successful update records the dtype actually passed to the quantizer. Shape broadcasting is unsupported because MXFP8 scale storage has fixed geometry.

copy_( value: torch.Tensor, ) → core.inference.quantization.mxfp8_tensor.MXFP8Tensor#

Tensor-compatible alias for :meth:quantize_ used by generic writeback.

Unlike torch.Tensor.copy_, this method requires an exact shape and does not accept non_blocking; quantization and storage updates are ordered on the current CUDA stream. The alias lets external conversion integrations such as Megatron Bridge treat plain and MXFP8 destinations uniformly while quantize_ remains available to existing callers.

classmethod from_bf16( x: torch.Tensor, group_size: int = 32, backend: core.inference.quantization.mxfp8_tensor.MXFP8Backend = 'flashinfer', )#

Quantize a floating-point CUDA tensor to MXFP8.

The historical method name is retained for compatibility. The Triton backend accepts BF16, FP16, and FP32 inputs; FlashInfer accepts BF16 and FP16 inputs.

Parameters:

x – [M, K] floating-point tensor on CUDA.
group_size – MXFP8 group size (default 32).
backend – ‘triton’ (fused quantize + swizzle Triton kernel) or ‘flashinfer’ (single fused FlashInfer CUDA kernel).

core.inference.quantization.mxfp8_tensor#