[Quantization][bugfix] Correct E8M0 NaN-sentinel detection in e8m0_to_f32 (#25519)

2026-07-09 15:02:53 -07:00
parent 966350408e
commit 48d98b7c68
2 changed files with 80 additions and 10 deletions
@@ -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()