[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

Fixes MXFP8 dequantization of UE8M0 edge codes.

• Expands code 0 to the FP32 subnormal value 2^-127 instead of zero.
• Expands code 255 to NaN instead of positive infinity.
• Adds a PyTorch regression test covering both edge codes.

Confidence Score: 5/5

The PR appears safe to merge with focused regression coverage for the corrected UE8M0 edge-code behavior.

The new branches return the required FP32 subnormal and NaN representations, and existing MXFP8 dequantization callers consume those values as block scales without exposing a conflicting contract.

Important Files Changed

Filename	Overview
transformer_engine/common/util/ptx.cuh	Correctly special-cases UE8M0 codes 0 and 255 consistently with the format semantics and the neighboring reciprocal implementation.
tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py	Adds focused regression coverage for both corrected dequantization cases and is discovered by the MXFP8 CI test invocation.

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