diff --git a/mlx/backend/cpu/unary_ops.h b/mlx/backend/cpu/unary_ops.h index f441e88bd5..e2dd7413f7 100644 --- a/mlx/backend/cpu/unary_ops.h +++ b/mlx/backend/cpu/unary_ops.h @@ -154,19 +154,20 @@ struct ToFP8 { struct FromFP8 { template Simd operator()(Simd x) { - auto v = Simd(x & 127) << 7; + auto v = Simd(x & 127); + 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) { - 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; } } - 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 42c5ec128a..00efdeb792 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -30,11 +30,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + 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; } operator bfloat16_t() thread { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 90a6501d97..75cfe6667a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -4258,6 +4258,17 @@ 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)) + # 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.assertTrue(np.isnan(decoded_np[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)