[CPU] Fix FP4 tie rounding on CPU

mlx/ops.cpp

@@ -5328,27 +5328,31 @@ std::vector<array> fp_quantize(
       scales = astype(add(scales, array(127, int32), s), uint8, s);
     }
     if (bits == 4) {
-      auto lut = array({
-          +0.0f,
-          +0.5f,
-          +1.0f,
-          +1.5f,
-          +2.0f,
-          +3.0f,
-          +4.0f,
-          +6.0f,
-          -0.0f,
-          -0.5f,
-          -1.0f,
-          -1.5f,
-          -2.0f,
-          -3.0f,
-          -4.0f,
-          -6.0f,
-      });
-      lut = astype(lut, w.dtype(), s);
-      wq = argmin(
-          abs(subtract(expand_dims(wq, -1, s), lut, s), s), -1, false, s);
+      // Round to nearest E2M1 with ties to even, same as fp4_e2m1 in Metal
+      constexpr std::pair<float, bool> bounds[] = {
+          {0.25f, false},
+          {0.75f, true},
+          {1.25f, false},
+          {1.75f, true},
+          {2.5f, false},
+          {3.5f, true},
+          {5.0f, false},
+      };
+      auto absw = abs(wq, s);
+      auto code = array(0, uint32);
+      for (auto [b, inclusive] : bounds) {
+        auto bound = array(b, wq.dtype());
+        auto above =
+            inclusive ? greater_equal(absw, bound, s) : greater(absw, bound, s);
+        code = add(code, astype(above, uint32, s), s);
+      }
+      // Read the sign bit so that -0 also gets the sign
+      auto int_type = size_of(wq.dtype()) == 2
+          ? int16
+          : (size_of(wq.dtype()) == 4 ? int32 : int64);
+      auto negative = less(view(wq, int_type, s), array(0, int_type), s);
+      code = where(negative, add(code, array(8, uint32), s), code, s);
+      wq = where(isnan(wq, s), array(7, uint32), code, s);
       auto shifts = power(array(2, uint32), arange(0, 32, 4, uint32, s), s);
       wq = reshape(wq, {-1, 4, 8}, s);
       wq = sum(multiply(wq, shifts, s), -1, false, s);

python/tests/test_quantized.py

@@ -131,6 +131,30 @@ def test_mxfp4_quantize_dequantize(self):
         w_hat = mx.dequantize(w_q, scales, mode="mxfp4")
         self.assertTrue(mx.all(w_hat == 0))
 
+    def test_fp4_quantize_ties_to_even(self):
+        vals = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 0.0]
+        vals = vals + [-v for v in vals]
+        expected = mx.array(
+            [0, 2, 2, 4, 4, 6, 6, 0, 8, 10, 10, 12, 12, 14, 14, 8], dtype=mx.uint32
+        )
+        shifts = mx.arange(0, 32, 4, dtype=mx.uint32)
+
+        devices = [mx.cpu]
+        if mx.is_available(mx.gpu):
+            devices.append(mx.gpu)
+        for mode, gs in [("mxfp4", 32), ("nvfp4", 16)]:
+            # Each group has 6.0, so the scale is 1
+            w = mx.zeros((len(vals), gs))
+            w[:, 0] = mx.array(vals)
+            w[:, 1] = 6.0
+            for dtype, device in product(
+                [mx.float32, mx.float16, mx.bfloat16], devices
+            ):
+                with self.subTest(mode=mode, dtype=dtype, device=device):
+                    w_q, _ = mx.quantize(w.astype(dtype), mode=mode, stream=device)
+                    codes = (w_q[:, 0:1] >> shifts)[:, 0] & 0xF
+                    self.assertTrue(mx.array_equal(codes, expected))
+
     def test_mxfp8_quantize_dequantize(self):
         w = 2 * mx.random.uniform(shape=(512, 32)) - 1
         w = w.astype(mx.bfloat16)