[common] Fix UE8M0 code 0 (2^-127) and code 255 (NaN) expansion in ptx::exp2f

[xiuhu17]

Description

ptx::exp2f(e8m0_t) builds the scale by shifting the biased exponent into the fp32 exponent field. Two UE8M0 codes cannot be built that way:

• code 0 is 2^-127 — an fp32 subnormal (0x00400000), so the shift yields +0.0. MXFP8 software dequantize (dequantize_mxfp8.cuh, group_dequantize_mxfp8.cuh) then zeroes every 1x32 block whose scale byte is 0.
• code 255 is NaN — the shift yields +Inf.

exp2f_rcp next to it already special-cases both; this applies the same handling to exp2f.

Repro:

import torch
from transformer_engine.pytorch import MXFP8Quantizer
import transformer_engine.pytorch as te

q = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, columnwise=False)
qx = q(torch.randn(32, 64, dtype=torch.bfloat16, device="cuda"))
qx._rowwise_data.view(torch.uint8)[0, :32] = 56   # E4M3 1.0
qx._rowwise_scale_inv.view(torch.uint8)[0, 0] = 0  # UE8M0 code 0 = 2^-127
print(qx.dequantize(dtype=torch.float32)[0, :4])
# before: tensor([0., 0., 0., 0.])
# after:  tensor([5.8775e-39, 5.8775e-39, 5.8775e-39, 5.8775e-39])  (= 2^-127)

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

Corrects UE8M0 extreme-code expansion and adds regression coverage.

• Expands code 0 to the FP32 subnormal value (2^{-127}) instead of zero.
• Expands code 255 to NaN instead of positive infinity.
• Tests MXFP8 dequantization with both extreme scale encodings.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/util/ptx.cuh	Adds explicit UE8M0 handling for the subnormal code 0 and NaN code 255 in the shared device expansion helper.
tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py	Adds an automatically discovered PyTorch regression test covering both corrected MXFP8 dequantization cases.

_{Reviews (2): Last reviewed commit: "Fix UE8M0 code 0 and 255 expansion in pt..." | Re-trigger Greptile}

[TristonC]

@ptrendx Anyone from your team can help review this?

[ptrendx]

/te-ci

[ptrendx]

@xiuhu17 Please sign your commit - please refer to https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work

Thanks!

[xiuhu17]

@xiuhu17 Please sign your commit - please refer to https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work

Thanks!

Fixed, thanks for reminding

[xiuhu17]

Hi @ptrendx, would you like to run the ci for this one? Thanks!

[ptrendx]

Actually, for this test to run you need to add it to the qa/L0_pytorch_unittest/test.sh

[xiuhu17]

Actually, for this test to run you need to add it to the qa/L0_pytorch_unittest/test.sh

Got it, my bad

[xiuhu17]

Actually, for this test to run you need to add it to the qa/L0_pytorch_unittest/test.sh

Hi @ptrendx, I think this one may already be covered. qa/L0_pytorch_unittest/test.sh:36 runs the whole directory rather than individual files:

python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8.xml $TE_PATH/tests/pytorch/mxfp8 || test_fail "test_mxfp8"