diff --git a/mlx/backend/cpu/unary_ops.h b/mlx/backend/cpu/unary_ops.h index f441e88bd5..1768a71f85 100644 --- a/mlx/backend/cpu/unary_ops.h +++ b/mlx/backend/cpu/unary_ops.h @@ -5,6 +5,7 @@ #include #include #include +#include #include "mlx/backend/cpu/simd/simd.h" @@ -165,6 +166,13 @@ struct FromFP8 { out[i] = converted * 256.0; } } + // 0x7f/0xff are e4m3's only NaN encodings. Shifted into float16 they land + // on a finite exponent (15, not 31), so the reinterpret above decodes them + // as 480. + out = select( + Simd((x & 127) == 127), + Simd(std::numeric_limits::quiet_NaN()), + out); auto sign = Simd(x & 128); return select(sign, -out, out); } diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 796dd21639..8a7642b918 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -33,6 +33,12 @@ struct fp8_e4m3 { uint16_t v = (bits & 127) << 7; half converted = as_type(v); converted *= 256.0; + // 0x7f/0xff are e4m3's only NaN encodings. Shifted into half they land on a + // finite exponent (15, not 31), so the reinterpret above decodes them as + // 480. + if ((bits & 127) == 127) { + converted = NAN; + } auto sign = bits & 128; return (sign ? -converted : converted); } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index f093891a15..4bc4596358 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3609,6 +3609,23 @@ 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)) + def test_from_fp8_nan(self): + # 0x7f/0xff are e4m3's only NaN encodings. Shifted into float16 they land on + # a finite exponent, so a plain reinterpret decodes them as +/-480. + nans = mx.array([0x7F, 0xFF], mx.uint8) + for dtype in [mx.float16, mx.bfloat16, mx.float32]: + for stream in [mx.cpu, mx.gpu]: + if stream == mx.gpu and not mx.metal.is_available(): + continue + out = mx.from_fp8(nans, dtype=dtype, stream=stream) + self.assertTrue( + mx.all(mx.isnan(out)).item(), msg=f"{dtype} {stream}: {out}" + ) + + # the largest finite magnitude is 448 (0x7e), and must stay finite + finite = mx.from_fp8(mx.array([0x7E, 0xFE], mx.uint8), dtype=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)