[MXLinear]Where is the operator call for implementing MXFP8 in NVD?

[LucaHW]

In the forward method of the MXLinear class, mx_mm.apply is called, although MXTensor.to_mx is also invoked. The following code implements the quantization processing of MXFP8：
scale_e8m0_biased, data_lp = to_mx(data_hp, elem_dtype, block_size, scaling_mode, is_swizzled_scales)

When examining the implementation of to_mx, I noticed that it does not call any CUDA-related low-precision operators; instead, it uses simulated low-precision implementations. What could be the reason for this? And where are the CUDA MXFP8 low-precision operators called? Thank you.

[danielvegamyhre]

emulated mode is just for CI / testing things if blackwell hardware is not available. the pytorch eager quantization ops are compiled to triton by inductor when you use torch.compile as intended. the mm op will dispatch to cublas mxfp8 gemm under the hood.

[LucaHW]

emulated mode is just for CI / testing things if blackwell hardware is not available. the pytorch eager quantization ops are compiled to triton by inductor when you use torch.compile as intended. the mm op will dispatch to cublas mxfp8 gemm under the hood.

Thank you for your answer. So if I don't use “torch.compile“, the process fails to enter the cuBLAS MXFP8 operator. Is this correct?
Can low-precision operators be invoked without enabling torch.compile?

[danielvegamyhre]

So if I don't use “torch.compile“, the process fails to enter the cuBLAS MXFP8 operator. Is this correct?

No, the torch.mm op with MXTensor args will dispatch to a mxfp8 cublas gemm, regardless of whether you're using eager or compile.

torch.compile speeds up the quantization ops by generating fused triton kernels for them. the same MXFP8 GEMM is used regardless though.

[LucaHW]

So if I don't use “torch.compile“, the process fails to enter the cuBLAS MXFP8 operator. Is this correct?

No, the torch.mm op with MXTensor args will dispatch to a mxfp8 cublas gemm, regardless of whether you're using eager or compile.

torch.compile speeds up the quantization ops by generating fused triton kernels for them. the same MXFP8 GEMM is used regardless though.

Thank you.
Therefore, ”torch._scaled_mm” takes MXFP8-type parameters for low-precision matrix multiplication computations, and this interface is the MXFP8 matrix multiplication operator provided by PyTorch. Prior to this, the quantization computation — specifically the conversion of high-precision BF16 format to MXFP8 format — is implemented via the interface "def to_mx". When torch.compile is not configured, each operator executes independently, resulting in relatively low performance. After torch.compile is enabled, these operators are fused into a single Triton kernel, which then improves execution efficiency.
Am I correct？
Is there no upper-level interface function for the to_mx method similar to torch._scaled_mm?