fix

mlx/backend/cpu/quantized.cpp

@@ -840,14 +840,15 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& w_pre = inputs[1];
   auto& scales_pre = inputs[2];
 
-  std::vector<array> temps;
+  auto& encoder = cpu::get_command_encoder(stream());
-  auto ensure_row_contiguous = [s = stream(), &temps](const array& arr) {
+  auto ensure_row_contiguous = [s = stream(), &encoder](const array& arr) {
     if (arr.flags().row_contiguous) {
       return arr;
     } else {
-      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {});
-      copy_cpu(arr, temps.back(), CopyType::General, s);
+      copy_cpu(arr, arr_cpy, CopyType::General, s);
-      return temps.back();
+      encoder.add_temporary(arr_cpy);
+      return arr_cpy;
     }
   };
 
@@ -857,8 +858,6 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   out.set_data(allocator::malloc(out.nbytes()));
 
-  auto& encoder = cpu::get_command_encoder(stream());
-  encoder.add_temporaries(std::move(temps));
   encoder.set_input_array(x);
   encoder.set_input_array(w);
   encoder.set_input_array(scales);
@@ -894,17 +893,18 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& lhs_indices = inputs[inputs.size() - 2];
   auto& rhs_indices = inputs[inputs.size() - 1];
 
-  std::vector<array> temps;
+  auto& encoder = cpu::get_command_encoder(stream());
   auto ensure_row_contiguous_last_dims = [s = stream(),
-                                          &temps](const array& arr) {
+                                          &encoder](const array& arr) {
     auto stride_0 = arr.strides()[arr.ndim() - 2];
     auto stride_1 = arr.strides()[arr.ndim() - 1];
     if (stride_0 == arr.shape(-1) && stride_1 == 1) {
       return arr;
     } else {
-      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {});
-      copy_cpu(arr, temps.back(), CopyType::General, s);
+      copy_cpu(arr, arr_cpy, CopyType::General, s);
-      return temps.back();
+      encoder.add_temporary(arr_cpy);
+      return arr_cpy;
     }
   };
 
@@ -914,8 +914,6 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
 
   out.set_data(allocator::malloc(out.nbytes()));
 
-  auto& encoder = cpu::get_command_encoder(stream());
-  encoder.add_temporaries(std::move(temps));
   encoder.set_input_array(x);
   encoder.set_input_array(w);
   encoder.set_input_array(scales);

python/src/ops.cpp

@@ -4206,7 +4206,8 @@ void init_ops(nb::module_& m) {
           ``quantize`` currently only supports 2D inputs with the second
           dimension divisible by ``group_size``
 
-        The supported quantization modes are described in more detail below.
+        The supported quantization modes are ``"affine"`` and ``"mxfp4"``. They
+        are described in more detail below.
 
         Args:
           w (array): Matrix to be quantized
@@ -4221,13 +4222,12 @@ void init_ops(nb::module_& m) {
 
           * w_q (array): The quantized version of ``w``
           * scales (array): The quantization scales
-          * biases (array): The quantization biases (returned for `mode=="affine"`).
+          * biases (array): The quantization biases (returned for ``mode=="affine"``).
 
         Notes:
-          The currently supported quantization mode is `"affine"`.
+          The ``affine`` mode quantizes groups of :math:`g` consecutive
-          Formally, for a group of :math:`g` consecutive elements :math:`w_1` to
+          elements in a row of ``w``. For each group the quantized
-          :math:`w_g` in a row of ``w`` we compute the quantized representation
+          representation of each element :math:`\hat{w_i}` is computed as follows:
-          of each element :math:`\hat{w_i}` as follows
 
           .. math::
 
@@ -4244,9 +4244,16 @@ void init_ops(nb::module_& m) {
           unsigned 32 bit integer where the 1st element occupies the 4 least
           significant bits, the 2nd bits 4-7 etc.
 
-          In order to be able to dequantize the elements of ``w`` we also need to
+          To dequantize the elements of ``w``, we also save :math:`s` and
-          save :math:`s` and :math:`\beta` which are the returned ``scales`` and
+          :math:`\beta` which are the returned ``scales`` and
-        ``biases`` respectively.
+          ``biases`` respectively.
+
+          The ``mxfp4`` mode similarly quantizes groups of :math:`g` elements
+          of ``w``. For ``mxfp4`` the group size must be ``32``. The elements
+          are quantized to 4-bit precision floating-point values (E2M1) with a
+          shared 8-bit scale per group. Unlike ``affine`` quantization,
+          ``mxfp4`` does not have a bias value. More details on the format can
+          be found in the `specification <https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf>`_.
       )pbdoc");
   m.def(
       "dequantize",
@@ -4264,11 +4271,9 @@ void init_ops(nb::module_& m) {
       R"pbdoc(
         Dequantize the matrix ``w`` using quantization parameters.
 
-        The supported quantization modes are described in more detail below.
-
         Args:
-          w (array): Matrix to be quantized
+          w (array): Matrix to be dequantized
-          scales (array): The scales to use per ``group_size`` elements of ``w``
+          scales (array): The scales to use per ``group_size`` elements of ``w``.
           biases (array, optional): The biases to use per ``group_size``
              elements of ``w``. Default: ``None``.
           group_size (int, optional): The size of the group in ``w`` that shares a
@@ -4281,10 +4286,11 @@ void init_ops(nb::module_& m) {
           array: The dequantized version of ``w``
 
         Notes:
-          The currently supported quantization mode is `"affine"`.
+          The currently supported quantization modes are ``"affine"`` and ``mxfp4``.
-          Formally, given the notation in :func:`quantize`, we compute
+
-          :math:`w_i` from :math:`\hat{w_i}` and corresponding :math:`s` and
+          For ``affine`` quantization, given the notation in :func:`quantize`,
-          :math:`\beta` as follows
+          we compute :math:`w_i` from :math:`\hat{w_i}` and corresponding :math:`s`
+          and :math:`\beta` as follows
 
           .. math::
 

tests/ops_tests.cpp

@@ -2996,7 +2996,10 @@ TEST_CASE("test quantize dequantize") {
 
   for (int i = 2; i <= 8; i *= 2) {
     int el_per_int = 32 / i;
-    auto [x_q, scales, biases] = quantize(x, 128, i);
+    auto res = quantize(x, 128, i);
+    auto x_q = res[0];
+    auto scales = res[1];
+    auto biases = res[2];
     CHECK_EQ(x_q.shape(), Shape{128, 512 / el_per_int});
     CHECK_EQ(scales.shape(), Shape{128, 4});
     CHECK_EQ(biases.shape(), Shape{128, 4});