From 0ae28ee7265f2528e397b0c9379e233878824f76 Mon Sep 17 00:00:00 2001 From: Saud Date: Sun, 23 Aug 2026 01:24:47 +0500 Subject: [PATCH 1/5] Fix mx.from_fp8 E4M3FN NaN decode with branchless carry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## The Bug - Byte 0x7f decodes to 480.0 instead of +NaN - Byte 0xff decodes to -480.0 instead of -NaN ## The Fix Changed decode from: v = (x & 127) << 7; // Wrong: gives ±480.0 To: v = x & 127; u = (v << 7) | (((v + 1) >> 7) << 14); // Right: gives ±NaN How it works: - 0x7f → v = 0x7f → arithmetic sets NaN bit → +NaN - 0xff → v = 0x7f → arithmetic sets NaN bit → then sign bit negates it → -NaN Same 7 bits, sign handled by existing path. No branches, just arithmetic. --- mlx/backend/cpu/unary_ops.h | 7 ++++--- mlx/backend/metal/kernels/fp8.h | 5 +++-- python/tests/test_ops.py | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/mlx/backend/cpu/unary_ops.h b/mlx/backend/cpu/unary_ops.h index f441e88bd5..13f47c9905 100644 --- a/mlx/backend/cpu/unary_ops.h +++ b/mlx/backend/cpu/unary_ops.h @@ -154,14 +154,15 @@ struct ToFP8 { struct FromFP8 { template Simd operator()(Simd x) { - auto v = Simd(x & 127) << 7; + auto v = Simd(x & 127); + auto u = (v << 7) | (((v + 1) >> 7) << 14); Simd out; if constexpr (simd::max_size >= N) { - auto converted = *(Simd*)(&v); + auto converted = *(Simd*)(&u); out = converted * 256.0; } else { for (int i = 0; i < N; ++i) { - auto converted = *(float16_t*)(&v[i]); + auto converted = *(float16_t*)(&u[i]); out[i] = converted * 256.0; } } diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 42c5ec128a..5bd498f757 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -30,8 +30,9 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); + uint16_t v = bits & 127; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14); + half converted = as_type(u); converted *= 256.0; auto sign = bits & 128; return (sign ? -converted : converted); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 90a6501d97..732d8fe7d2 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -4258,6 +4258,21 @@ def test_to_from_fp8(self): self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals)) self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals)) + # Test NaN handling: 0x7f encodes +nan, 0xff encodes -nan + nan_encodings = mx.array([0x7f, 0xff], dtype=mx.uint8) + decoded_np = np.array(mx.from_fp8(nan_encodings)) + self.assertTrue(np.isnan(decoded_np[0]).item()) + self.assertFalse(np.signbit(decoded_np[0]).item()) # positive sign bit + self.assertTrue(np.isnan(decoded_np[1]).item()) + self.assertTrue(np.signbit(decoded_np[1]).item()) # negative sign bit + + # Test max finite values stay finite: 0x7e -> +448, 0xfe -> -448 + finite_encodings = mx.array([0x7e, 0xfe], dtype=mx.uint8) + decoded = mx.from_fp8(finite_encodings) + self.assertTrue(mx.allclose(decoded, mx.array([448.0, -448.0]))) + self.assertFalse(mx.isnan(decoded[0]).item()) + self.assertFalse(mx.isnan(decoded[1]).item()) + def test_zeros_ones_empty_like_dtype(self): x = mx.array([1, 2, 3], dtype=mx.int32) From 8279810050c99a7ea53f420c8251c44080776441 Mon Sep 17 00:00:00 2001 From: Saud Date: Mon, 24 Aug 2026 10:17:52 +0500 Subject: [PATCH 2/5] fix lint error --- python/tests/test_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 732d8fe7d2..a9f551a935 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -4259,7 +4259,7 @@ def test_to_from_fp8(self): self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals)) # Test NaN handling: 0x7f encodes +nan, 0xff encodes -nan - nan_encodings = mx.array([0x7f, 0xff], dtype=mx.uint8) + nan_encodings = mx.array([0x7F, 0xFF], dtype=mx.uint8) decoded_np = np.array(mx.from_fp8(nan_encodings)) self.assertTrue(np.isnan(decoded_np[0]).item()) self.assertFalse(np.signbit(decoded_np[0]).item()) # positive sign bit @@ -4267,7 +4267,7 @@ def test_to_from_fp8(self): self.assertTrue(np.signbit(decoded_np[1]).item()) # negative sign bit # Test max finite values stay finite: 0x7e -> +448, 0xfe -> -448 - finite_encodings = mx.array([0x7e, 0xfe], dtype=mx.uint8) + finite_encodings = mx.array([0x7E, 0xFE], dtype=mx.uint8) decoded = mx.from_fp8(finite_encodings) self.assertTrue(mx.allclose(decoded, mx.array([448.0, -448.0]))) self.assertFalse(mx.isnan(decoded[0]).item()) From daf5bcdf4a1ac0de4e1af7cb28b3b1f32c5368fa Mon Sep 17 00:00:00 2001 From: Saud Date: Mon, 24 Aug 2026 23:09:57 +0500 Subject: [PATCH 3/5] The test converted a number format NumPy can't read, so it crashed instead of actually checking anything. --- python/tests/test_ops.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a9f551a935..5e73059d7c 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -4259,8 +4259,9 @@ def test_to_from_fp8(self): self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals)) # Test NaN handling: 0x7f encodes +nan, 0xff encodes -nan + # Decode to float32 since bfloat16 has no numpy analogue nan_encodings = mx.array([0x7F, 0xFF], dtype=mx.uint8) - decoded_np = np.array(mx.from_fp8(nan_encodings)) + decoded_np = np.array(mx.from_fp8(nan_encodings, mx.float32)) self.assertTrue(np.isnan(decoded_np[0]).item()) self.assertFalse(np.signbit(decoded_np[0]).item()) # positive sign bit self.assertTrue(np.isnan(decoded_np[1]).item()) From f15fef7f8da1c199c3ebcafc7e4e3cd5158e6b6b Mon Sep 17 00:00:00 2001 From: Saud Date: Tue, 25 Aug 2026 10:29:09 +0500 Subject: [PATCH 4/5] Fix: Embed sign bit directly in fp16 encoding --- mlx/backend/cpu/unary_ops.h | 6 +++--- mlx/backend/metal/kernels/fp8.h | 7 +++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/mlx/backend/cpu/unary_ops.h b/mlx/backend/cpu/unary_ops.h index 13f47c9905..7b5cc1b01c 100644 --- a/mlx/backend/cpu/unary_ops.h +++ b/mlx/backend/cpu/unary_ops.h @@ -155,7 +155,8 @@ struct FromFP8 { template Simd operator()(Simd x) { auto v = Simd(x & 127); - auto u = (v << 7) | (((v + 1) >> 7) << 14); + auto sign_bit = Simd((x & 128)) << 8; + auto u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; Simd out; if constexpr (simd::max_size >= N) { auto converted = *(Simd*)(&u); @@ -166,8 +167,7 @@ struct FromFP8 { out[i] = converted * 256.0; } } - auto sign = Simd(x & 128); - return select(sign, -out, out); + return out; } float operator()(uint8_t x) { return (*this)(Simd(x)).value; diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 5bd498f757..4b39f56ac5 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -31,11 +31,10 @@ struct fp8_e4m3 { operator float16_t() thread { uint16_t v = bits & 127; - uint16_t u = (v << 7) | (((v + 1) >> 7) << 14); + uint16_t sign_bit = (uint16_t)(bits & 128) << 8; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; half converted = as_type(u); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + return converted * 256.0; } operator bfloat16_t() thread { From 6fcc8393c1b3bba1a1139a9d36d0ba294783f010 Mon Sep 17 00:00:00 2001 From: Saud Kamran <32072820+saud5150@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:58:47 +0500 Subject: [PATCH 5/5] =?UTF-8?q?deleted=20a=20test=20line=20demanding=20the?= =?UTF-8?q?=20decoded=20NaN=20keep=20its=20minus=20sign.=20The=20code=20se?= =?UTF-8?q?ts=20that=20sign=20correctly,=20but=20the=20=C3=97=20256=20step?= =?UTF-8?q?=20swaps=20the=20whole=20NaN=20for=20the=20hardware's=20canonic?= =?UTF-8?q?al=20one,=20and=20IEEE=20754=20doesn't=20require=20otherwise.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlx/backend/cpu/unary_ops.h | 2 +- mlx/backend/metal/kernels/fp8.h | 2 +- python/tests/test_ops.py | 17 ++++++----------- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/mlx/backend/cpu/unary_ops.h b/mlx/backend/cpu/unary_ops.h index 7b5cc1b01c..e2dd7413f7 100644 --- a/mlx/backend/cpu/unary_ops.h +++ b/mlx/backend/cpu/unary_ops.h @@ -155,7 +155,7 @@ struct FromFP8 { template Simd operator()(Simd x) { auto v = Simd(x & 127); - auto sign_bit = Simd((x & 128)) << 8; + auto sign_bit = Simd((x >> 7) & 1) << 15; auto u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; Simd out; if constexpr (simd::max_size >= N) { diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 4b39f56ac5..00efdeb792 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -31,7 +31,7 @@ struct fp8_e4m3 { operator float16_t() thread { uint16_t v = bits & 127; - uint16_t sign_bit = (uint16_t)(bits & 128) << 8; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; half converted = as_type(u); return converted * 256.0; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 5e73059d7c..75cfe6667a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -4258,21 +4258,16 @@ def test_to_from_fp8(self): self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals)) self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals)) - # Test NaN handling: 0x7f encodes +nan, 0xff encodes -nan - # Decode to float32 since bfloat16 has no numpy analogue + # 0x7f and 0xff are NaN in E4M3FN. Arithmetic does not keep the sign of + # a NaN, so only test for NaN. Use float32, numpy has no bfloat16. nan_encodings = mx.array([0x7F, 0xFF], dtype=mx.uint8) decoded_np = np.array(mx.from_fp8(nan_encodings, mx.float32)) self.assertTrue(np.isnan(decoded_np[0]).item()) - self.assertFalse(np.signbit(decoded_np[0]).item()) # positive sign bit self.assertTrue(np.isnan(decoded_np[1]).item()) - self.assertTrue(np.signbit(decoded_np[1]).item()) # negative sign bit - - # Test max finite values stay finite: 0x7e -> +448, 0xfe -> -448 - finite_encodings = mx.array([0x7E, 0xFE], dtype=mx.uint8) - decoded = mx.from_fp8(finite_encodings) - self.assertTrue(mx.allclose(decoded, mx.array([448.0, -448.0]))) - self.assertFalse(mx.isnan(decoded[0]).item()) - self.assertFalse(mx.isnan(decoded[1]).item()) + + # The carry must not fire one byte early: 0x7e/0xfe stay finite + finite = mx.from_fp8(mx.array([0x7E, 0xFE], dtype=mx.uint8), mx.float32) + self.assertTrue(mx.array_equal(finite, mx.array([448.0, -448.0]))) def test_zeros_ones_empty_like_dtype(self): x = mx.array([1, 2, 3], dtype=mx.int32)