MXFP4 Tensor Core GEMM support in GPT-OSS for Blackwell GPUs
Author: TheTinyTeddy
Pull request: no

Hi there,

I was looking at the code (in `triton_kernels/matmul_ogs_details/_matmul_ogs.py`) and found that even when using Blackwell GPU the Triton kernel that implements 
`acc = tl.dot_scaled(x, x_scales, x_format, w, w_scales, w_format, acc=acc, fast_math=True)`

is actually doing BF16 GEMM rather than MXFP4 GEMM, and x_scales is actually None, i.e. no quantization for the activation.

Therefore, I was wondering will there be support for native MXFP4 GEMM in Transformers?

Many thanks


Comments:
cc @arthurzucker @sunmarc but I thought our kernels supported MXFP4 in the forward direction at least?

> cc [@ArthurZucker](https://github.com/ArthurZucker) [@SunMarc](https://github.com/SunMarc) but I thought our kernels supported MXFP4 in the forward direction at least?

At the first glance of the code I was thinking the same too, but I then debug through the code and inspect the compiled kernel's ptx and found the mma instruction is actually doing bf16/bf16->fp32 GEMM

If i'm not wrong, for gpt-oss, we don't quantize the activations at all. So this is normal that we are doing BF16 GEMM. Only the weights are quantized `w, w_scales, w_format ` and we dequantize them during the matmult to perform bf16 gemm. 

> cc @ArthurZucker @SunMarc but I thought our kernels supported MXFP4 in the forward direction at least?

It supports `bf16 x mxfp4` op (activation x weights) thanks to the triton kernels and this is how OAI suggest to perform inference on it. 

> If i'm not wrong, for gpt-oss, we don't quantize the activations at all. So this is normal that we are doing BF16 GEMM. Only the weights are quantized `w, w_scales, w_format ` and we dequantize them during the matmult to perform bf16 gemm.
> 
> >

So GPT-OSS wouldn't be able to utilize the Blackwell's mxfp4 tensor core. I guess another reason is fp4 quantization of the activation (especially without AWQ/SmoothQuant-like smoothing) can introduce significant error?

> So GPT-OSS wouldn't be able to utilize the Blackwell's mxfp4 tensor core. I guess another reason is fp4 quantization of the activation (especially without AWQ/SmoothQuant-like smoothing) can introduce significant error?

Most likely. But note that some framework like vllm say that they have support for Blackwell's mxfp4 tensor core. So I think the degradation is not as bad as you would expect. If you check their implementation, they uses flashinfer kernels and they quantize the activation in mxfp8 https://github.com/vllm-project/vllm/blob/ef7eefe17a7dc212ddb8a8aabd7760218a10e25e/vllm/model_executor/layers/quantization/mxfp4.py#L895



> > So GPT-OSS wouldn't be able to utilize the Blackwell's mxfp4 tensor core. I guess another reason is fp4 quantization of the activation (especially without AWQ/SmoothQuant-like smoothing) can introduce significant error?
> 
> Most likely. But note that some framework like vllm say that they have support for Blackwell's mxfp4 tensor core. So I think the degradation is not as bad as you would expect. If you check their implementation, they uses flashinfer kernels and they quantize the activation in mxfp8 https://github.com/vllm-project/vllm/blob/ef7eefe17a7dc212ddb8a8aabd7760218a10e25e/vllm/model_executor/layers/quantization/mxfp4.py#L895

I think quantizing the activation to mxfp8 can work well, but it can only leverage the mxfp8 tensor core. However, there maybe difficulty for quantizing the activation using mxfp4. There is a factor of 2 or 3 times the FLOPs between mxfp8 and mxfp4 tensor core in Blackwell, so the performance will be lower.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.