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

tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch import MXFP8Quantizer
+
+recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_dequantize_extreme_e8m0_scale_codes() -> None:
+    """UE8M0 scale code 0 is 2^-127 and code 255 is NaN, not 0.0 and +Inf."""
+    quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, columnwise=False)
+    x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+    qx = quantizer(x)
+    data = qx._rowwise_data.view(torch.uint8)
+    scales = qx._rowwise_scale_inv.view(torch.uint8)
+    data[0, :32] = 56
+    scales[0, 0] = 0
+    data[1, :32] = 56
+    scales[1, 0] = 255
+    y = qx.dequantize(dtype=torch.float32)
+    assert (y[0, :32] == 2.0**-127).all()
+    assert torch.isnan(y[1, :32]).all()

transformer_engine/common/util/ptx.cuh

@@ -380,6 +380,10 @@ __device__ __forceinline__ bf16 exp2f_rcp<bf16>(e8m0_t biased_exp) {
 }
 
 __device__ __forceinline__ float exp2f(e8m0_t biased_exp) {
+  // Handle the special case of NaN.
+  if (biased_exp == 255) return __int_as_float(0x7fffffff);
+  // 2^-127 is subnormal, so it cannot be built by shifting into the exponent field.
+  if (biased_exp == 0) return __int_as_float(0x00400000);
   return __int_as_float(biased_exp << FP32_MANTISSA_BITS);
 }