From 48d98b7c6885b141d652bf32bbeed61dcb74c8a4 Mon Sep 17 00:00:00 2001
From: Spandan Tiwari <23646532+spandantiwari@users.noreply.github.com>
Date: Thu, 9 Jul 2026 15:02:53 -0700
Subject: [PATCH] [Quantization][bugfix] Correct E8M0 NaN-sentinel detection in
 e8m0_to_f32 (#25519)

---
 .../srt/layers/quantization/quark/utils.py    | 16 ++--
 .../layers/quantization/test_quark_utils.py   | 74 +++++++++++++++++++
 2 files changed, 80 insertions(+), 10 deletions(-)
 create mode 100644 test/registered/unit/layers/quantization/test_quark_utils.py

diff --git a/python/sglang/srt/layers/quantization/quark/utils.py b/python/sglang/srt/layers/quantization/quark/utils.py
index 7fd865233..4c2d162f0 100644
--- a/python/sglang/srt/layers/quantization/quark/utils.py
+++ b/python/sglang/srt/layers/quantization/quark/utils.py
@@ -13,7 +13,7 @@ except ImportError:
 
     def raise_aiter_import_error(*args, **kwargs):
         raise ImportError(
-            "Failed to import aiter. " "Make sure AITER is installed and accessible."
+            "Failed to import aiter. Make sure AITER is installed and accessible."
         )
 
     dynamic_mxfp4_quant = raise_aiter_import_error
@@ -161,16 +161,12 @@ def mxfp4_to_f32(x, is_3d):
 
 
 def e8m0_to_f32(x):
-    # Convert the input tensor `x` (assumed to be in e8m0 format) to float32.
-    # e8m0 is a custom 8-bit floating point format with 8 bits for exponent, 0 for mantissa.
-    # This means the value is essentially 2^(exponent - 127), similar to how IEEE-754 stores floats.
-
-    # Convert x to float32 for computation, and compute the power of 2 by subtracting the bias (127).
+    # Per OCP MX-format v1.0: encoded 0..254 -> 2^(x-127); encoded 255 -> NaN.
+    # Detect the sentinel on the raw integer encoding, not on the float result
+    # (in float32, 2^128 overflows to +inf, so the old `x_f32 == 128` predicate
+    # both missed x=255 and wrongly NaN'd legitimate scale 128.0 at x=134).
     x_f32 = 2 ** ((x.to(torch.float32)) - 127)
-
-    # If the exponent value was 255 (i.e., 2^(128)), this is a special case usually used to represent NaN or Inf.
-    # Since this custom format has no mantissa, treat 2^128 as NaN.
-    x_f32[x_f32 == 128] = float("nan")
+    x_f32[x == 255] = float("nan")
     return x_f32
 
 
diff --git a/test/registered/unit/layers/quantization/test_quark_utils.py b/test/registered/unit/layers/quantization/test_quark_utils.py
new file mode 100644
index 000000000..fd67bde62
--- /dev/null
+++ b/test/registered/unit/layers/quantization/test_quark_utils.py
@@ -0,0 +1,74 @@
+"""Unit tests for sglang.srt.layers.quantization.quark.utils — CPU-only, no model loading."""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.quark.utils import e8m0_to_f32
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestE8M0ToF32(CustomTestCase):
+    """Cover OCP MX-format v1.0 e8m0 decoding:
+    encoded 0..254 -> 2^(x-127); encoded 255 -> NaN.
+    """
+
+    # ---- Bug-catchers: must FAIL on unfixed code ----------------------------
+
+    def test_scale_128_is_not_nan(self):
+        # Bug facet 1: legit scale 128.0 (x=134) was being poisoned to NaN.
+        x = torch.tensor([134], dtype=torch.uint8)
+        out = e8m0_to_f32(x)
+        self.assertEqual(out.item(), 128.0)
+        self.assertFalse(torch.isnan(out).any().item())
+
+    def test_nan_sentinel(self):
+        # Bug facet 2: x=255 is the OCP NaN sentinel; was passing through as +inf.
+        x = torch.tensor([255], dtype=torch.uint8)
+        self.assertTrue(torch.isnan(e8m0_to_f32(x)).all().item())
+
+    def test_only_255_is_nan(self):
+        # Exactly one of 0..255 should be NaN, and it must be index 255.
+        # Build the range in the default int dtype then cast — passing the
+        # uint8 dtype directly to `arange(0, 256, dtype=uint8)` raises on
+        # PyTorch versions that bounds-check the end value (256 is out of
+        # uint8 range).
+        x = torch.arange(256).to(torch.uint8)
+        out = e8m0_to_f32(x)
+        nan_idx = torch.isnan(out).nonzero().flatten().tolist()
+        self.assertEqual(nan_idx, [255])
+
+    def test_known_powers_of_two(self):
+        x = torch.tensor([0, 125, 126, 127, 128, 129, 134, 254], dtype=torch.uint8)
+        expected = torch.tensor(
+            [2.0**-127, 0.25, 0.5, 1.0, 2.0, 4.0, 128.0, 2.0**127],
+            dtype=torch.float32,
+        )
+        torch.testing.assert_close(e8m0_to_f32(x), expected)
+
+    # ---- Guardrails: pass on both buggy and fixed code ----------------------
+
+    def test_zero_exponent_is_one(self):
+        x = torch.tensor([127], dtype=torch.uint8)
+        self.assertEqual(e8m0_to_f32(x).item(), 1.0)
+
+    def test_shape_preserved(self):
+        x = torch.zeros((3, 4, 5), dtype=torch.uint8)
+        self.assertEqual(tuple(e8m0_to_f32(x).shape), (3, 4, 5))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "no GPU")
+    def test_cuda_parity(self):
+        x = torch.tensor([127, 134, 255], dtype=torch.uint8, device="cuda")
+        out = e8m0_to_f32(x)
+        self.assertEqual(out.device.type, "cuda")
+        self.assertEqual(out[0].item(), 1.0)
+        self.assertEqual(out[1].item(), 128.0)
+        self.assertTrue(torch.isnan(out[2]).item())
+
+
+if __name__ == "__main__":
+    unittest.main()