From 4e9e807d412bf5daefe367f57b2dd2948f014789 Mon Sep 17 00:00:00 2001
From: zhihaow6 <zhihaow6@illinois.edu>
Date: Sat, 25 Jul 2026 11:00:46 -0700
Subject: [PATCH] Fix UE8M0 code 0 and 255 expansion in ptx::exp2f

UE8M0 code 0 is 2^-127 and code 255 is NaN, but the exponent-field shift
produced +0.0 and +Inf, so MXFP8 software dequantize zeroed every 1x32
block whose scale byte was 0. Mirror the special cases already present in
exp2f_rcp and add a dequantize test with planted extreme scale codes.
---
 .../test_mxfp8_dequantize_extreme_scales.py   | 28 +++++++++++++++++++
 transformer_engine/common/util/ptx.cuh        |  4 +++
 2 files changed, 32 insertions(+)
 create mode 100644 tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py

diff --git a/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py
new file mode 100644
index 0000000000..64bfe6012e
--- /dev/null
+++ b/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()
diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh
index 2814aa3490..df21a0af56 100644
--- a/transformer_engine/common/util/ptx.cuh
+++ b/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);
 }