From 27e31ab2499fa74f4f7be4a477f0c33a5f7eef45 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 21 Aug 2025 15:54:30 -0700 Subject: [PATCH] fix --- mlx/backend/cpu/quantized.cpp | 26 +++++++++++------------ python/src/ops.cpp | 40 ++++++++++++++++++++--------------- tests/ops_tests.cpp | 5 ++++- 3 files changed, 39 insertions(+), 32 deletions(-) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index ce51cac3a..39914838b 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -840,14 +840,15 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { auto& w_pre = inputs[1]; auto& scales_pre = inputs[2]; - std::vector temps; - auto ensure_row_contiguous = [s = stream(), &temps](const array& arr) { + auto& encoder = cpu::get_command_encoder(stream()); + 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, {})); - copy_cpu(arr, temps.back(), CopyType::General, s); - return temps.back(); + auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {}); + copy_cpu(arr, arr_cpy, CopyType::General, s); + encoder.add_temporary(arr_cpy); + return arr_cpy; } }; @@ -857,8 +858,6 @@ void QuantizedMatmul::eval_cpu(const std::vector& 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& inputs, array& out) { auto& lhs_indices = inputs[inputs.size() - 2]; auto& rhs_indices = inputs[inputs.size() - 1]; - std::vector 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, {})); - copy_cpu(arr, temps.back(), CopyType::General, s); - return temps.back(); + auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {}); + copy_cpu(arr, arr_cpy, CopyType::General, s); + encoder.add_temporary(arr_cpy); + return arr_cpy; } }; @@ -914,8 +914,6 @@ void GatherQMM::eval_cpu(const std::vector& 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); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index cb0add614..3e7218ad5 100644 --- a/python/src/ops.cpp +++ b/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"`. - Formally, for a group of :math:`g` consecutive elements :math:`w_1` to - :math:`w_g` in a row of ``w`` we compute the quantized representation - of each element :math:`\hat{w_i}` as follows + The ``affine`` mode quantizes groups of :math:`g` consecutive + elements in a row of ``w``. For each group the quantized + representation of each element :math:`\hat{w_i}` is computed 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 - save :math:`s` and :math:`\beta` which are the returned ``scales`` and - ``biases`` respectively. + To dequantize the elements of ``w``, we also save :math:`s` and + :math:`\beta` which are the returned ``scales`` and + ``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 `_. )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 - scales (array): The scales to use per ``group_size`` elements of ``w`` + w (array): Matrix to be dequantized + 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"`. - Formally, given the notation in :func:`quantize`, we compute - :math:`w_i` from :math:`\hat{w_i}` and corresponding :math:`s` and - :math:`\beta` as follows + The currently supported quantization modes are ``"affine"`` and ``mxfp4``. + + For ``affine`` quantization, given the notation in :func:`quantize`, + we compute :math:`w_i` from :math:`\hat{w_i}` and corresponding :math:`s` + and :math:`\beta` as follows .. math:: diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 17207efd4..878c7101b 100644 --- a/tests/ops_tests.cpp +++ b/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});