zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-29 18:26:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Awni Hannun 2025-08-21 15:54:30 -07:00
parent 80868ee4fb
commit 27e31ab249
3 changed files with 39 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
mlx/backend/cpu/quantized.cpp
View File

@ -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 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<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, {}));
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<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);

40
python/src/ops.cpp
View File

@ -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 <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
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::

5
tests/ops_tests.cpp
View File

@ -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});