[Inquiry] Support status and roadmap for MXFP8 on SM120 (Blackwell)

[Mnb66]

Is your feature request related to a problem? Please describe.

I am currently exploring the performance benefits of the Blackwell architecture, specifically focusing on the SM120 compute capability. While recent releases mention MXFP8 support, there seems to be limited documentation or potentially missing features specifically for the SM120 target compared to SM100.

Describe the solution you'd like

I would appreciate clarification on:

• The official support status and feature roadmap for MXFP8 on SM120 devices.
• Whether there are plans to introduce SM120-optimized kernels specifically for MXFP8BlockScaling that differ from the current Float8BlockScaling execution path.

Describe alternatives you've considered

I have tested the Float8BlockScaling recipe, and it functions correctly on SM120. However, I believe that MXFP8BlockScaling might potentially offer better performance or more native hardware utilization on Blackwell arch.

Additional context

After reviewing the TransformerEngine source code, I observed that on SM120 (and generally for sm_arch() >= 100), the Float8BlockScaling implementation performs a conversion of 1D/2D block scaling into an MXFP8 layout with E8M0 scales. It then invokes the same MXFP8 GEMM path as MXFP8BlockScaling.

Since both recipes eventually converge on the same GEMM kernel execution path, I have a few technical questions regarding this design:

• What is the primary design purpose behind having Float8BlockScaling and MXFP8BlockScaling share the same kernel path on Blackwell?
• Does the conversion step in Float8BlockScaling introduce significant overhead compared to a native MXFP8BlockScaling approach?

[ptrendx]

Hi @Mnb66.

You are correct that currently the support for SM120 lags behind SM100/103. We are actively working with cuBLAS and cuDNN teams to bridge those gaps in functionality.
For MXFP8 specifically the main issue is lack of support for non-TN GEMM layouts in cuBLAS at the moment. It is not a fundamental limitation and we are working with cuBLAS on adding that support. This limitation means that, while the MXFP8 forward pass works on SM120, backward pass currently does not.
As for the differences between the block scaling and MXFP8 scaling - you are correct that the way block scaling is implemented for Blackwell in general (not only for SM120, but also for SM100) is emulated via the MXFP8 code path. The reason for that is that the MXFP8 is natively accelerated on Blackwell, while the Blockwise recipe is not. This emulation is possible because the granularity of the scaling factor in MXFP8 case is strictly higher than in the blockwise recipe, so by broadcasting the scaling factors we can get the same result. The Hopper way of doing blockwise scaling was not natively accelerated by hardware, which resulted in slowdown compared to the fully native GEMM.
The conversion step should not incur significant overhead compared to the native MXFP8 recipe, but that depends on the exact tensor shapes etc. in your case. The overhead comes from broadcasting the scaling factors, which constitute ~3% of the data. However, for smaller shapes the efficiency of that broadcast will be lower, making the cost relatively higher.

Does that answer your questions?