mlx_rs/ops/
quantization.rs

1use std::ffi::CStr;
2
3use mlx_internal_macros::generate_macro;
4
5use crate::{
6    error::Result,
7    utils::{guard::Guarded, VectorArray},
8    Array, Stream,
9};
10
11const DEFAULT_MODE: &CStr = c"affine";
12const DEFAULT_GROUP_SIZE: i32 = 64;
13const DEFAULT_BITS: i32 = 4;
14
15/// Helper to convert Option<i32> to mlx_optional_int
16fn optional_int(value: Option<i32>, default: i32) -> mlx_sys::mlx_optional_int {
17    mlx_sys::mlx_optional_int {
18        value: value.unwrap_or(default),
19        has_value: value.is_some(),
20    }
21}
22
23/// Helper to create a "no value" optional dtype
24fn optional_dtype_none() -> mlx_sys::mlx_optional_dtype {
25    mlx_sys::mlx_optional_dtype {
26        value: mlx_sys::mlx_dtype__MLX_FLOAT32, // default value, ignored when has_value is false
27        has_value: false,
28    }
29}
30
31/// Quantize the matrix `w` using `bits` bits per element.
32///
33/// Note, every `group_size` elements in a row of `w` are quantized together. Hence, number of
34/// columns of `w` should be divisible by `group_size`. In particular, the rows of `w` are divided
35/// into groups of size `group_size` which are quantized together.
36///
37/// > `quantized` currently only supports 2D inputs with dimensions which are multiples of 32
38///
39/// For details, please see [this
40/// documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.quantize.html)
41///
42/// # Params
43///
44/// - `w`: The input matrix
45/// - `group_size`: The size of the group in `w` that shares a scale and bias. (default: `64`)
46/// - `bits`: The number of bits occupied by each element of w in the returned quantized matrix.
47///   (default: 4)
48pub fn quantize(
49    w: impl AsRef<Array>,
50    group_size: impl Into<Option<i32>>,
51    bits: impl Into<Option<i32>>,
52) -> Result<(Array, Array, Array)> {
53    let stream = Stream::thread_local_or_default();
54    let group_size = optional_int(group_size.into(), DEFAULT_GROUP_SIZE);
55    let bits = optional_int(bits.into(), DEFAULT_BITS);
56    let global_scale = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
57
58    let result = VectorArray::try_from_op(|res| unsafe {
59        mlx_sys::mlx_quantize(
60            res,
61            w.as_ref().as_ptr(),
62            group_size,
63            bits,
64            DEFAULT_MODE.as_ptr(),
65            global_scale.as_ptr(),
66            stream.as_ref().as_ptr(),
67        )
68    })?;
69
70    let arrays: Vec<Array> = result.try_into_values()?;
71    if arrays.len() != 3 {
72        return Err(crate::error::Exception::custom(format!(
73            "Expected 3 arrays from quantize, got {}",
74            arrays.len()
75        )));
76    }
77    let mut iter = arrays.into_iter();
78    Ok((
79        iter.next().unwrap(),
80        iter.next().unwrap(),
81        iter.next().unwrap(),
82    ))
83}
84
85/// Compatibility shim for [`quantize`].
86#[generate_macro(customize(forwarding_shim = true))]
87#[deprecated(
88    since = "0.26.0",
89    note = "use `with_stream` or `with_device` around `quantize`"
90)]
91pub fn quantize_device(
92    w: impl AsRef<Array>,
93    #[optional] group_size: impl Into<Option<i32>>,
94    #[optional] bits: impl Into<Option<i32>>,
95    #[optional] stream: impl AsRef<Stream>,
96) -> Result<(Array, Array, Array)> {
97    crate::with_stream(stream.as_ref(), || quantize(w, group_size, bits))
98}
99
100/// Perform the matrix multiplication with the quantized matrix `w`. The quantization uses one
101/// floating point scale and bias per `group_size` of elements. Each element in `w` takes `bits`
102/// bits and is packed in an unsigned 32 bit integer.
103#[allow(clippy::too_many_arguments)]
104pub fn quantized_matmul<'a>(
105    x: impl AsRef<Array>,
106    w: impl AsRef<Array>,
107    scales: impl AsRef<Array>,
108    biases: impl Into<Option<&'a Array>>,
109    transpose: impl Into<Option<bool>>,
110    group_size: impl Into<Option<i32>>,
111    bits: impl Into<Option<i32>>,
112) -> Result<Array> {
113    let stream = Stream::thread_local_or_default();
114    let transpose = transpose.into().unwrap_or(false);
115    let group_size = optional_int(group_size.into(), DEFAULT_GROUP_SIZE);
116    let bits = optional_int(bits.into(), DEFAULT_BITS);
117
118    <Array as Guarded>::try_from_op(|res| unsafe {
119        mlx_sys::mlx_quantized_matmul(
120            res,
121            x.as_ref().as_ptr(),
122            w.as_ref().as_ptr(),
123            scales.as_ref().as_ptr(),
124            biases
125                .into()
126                .map(|a| a.as_ptr())
127                .unwrap_or(mlx_sys::mlx_array_new()),
128            transpose,
129            group_size,
130            bits,
131            DEFAULT_MODE.as_ptr(),
132            stream.as_ref().as_ptr(),
133        )
134    })
135}
136
137/// Compatibility shim for [`quantized_matmul`].
138#[allow(clippy::too_many_arguments)]
139#[generate_macro(customize(forwarding_shim = true))]
140#[deprecated(
141    since = "0.26.0",
142    note = "use `with_stream` or `with_device` around `quantized_matmul`"
143)]
144pub fn quantized_matmul_device<'a>(
145    x: impl AsRef<Array>,
146    w: impl AsRef<Array>,
147    scales: impl AsRef<Array>,
148    #[optional] biases: impl Into<Option<&'a Array>>,
149    #[optional] transpose: impl Into<Option<bool>>,
150    #[optional] group_size: impl Into<Option<i32>>,
151    #[optional] bits: impl Into<Option<i32>>,
152    #[optional] stream: impl AsRef<Stream>,
153) -> Result<Array> {
154    crate::with_stream(stream.as_ref(), || {
155        quantized_matmul(x, w, scales, biases, transpose, group_size, bits)
156    })
157}
158
159/// Dequantize the matrix `w` using the provided `scales` and `biases` and the `group_size` and
160/// `bits` configuration.
161///
162/// For details, please see [this
163/// documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.dequantize.html)
164pub fn dequantize<'a>(
165    w: impl AsRef<Array>,
166    scales: impl AsRef<Array>,
167    biases: impl Into<Option<&'a Array>>,
168    group_size: impl Into<Option<i32>>,
169    bits: impl Into<Option<i32>>,
170) -> Result<Array> {
171    let stream = Stream::thread_local_or_default();
172    let group_size = optional_int(group_size.into(), DEFAULT_GROUP_SIZE);
173    let bits = optional_int(bits.into(), DEFAULT_BITS);
174    let global_scale = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
175
176    <Array as Guarded>::try_from_op(|res| unsafe {
177        mlx_sys::mlx_dequantize(
178            res,
179            w.as_ref().as_ptr(),
180            scales.as_ref().as_ptr(),
181            biases
182                .into()
183                .map(|a| a.as_ptr())
184                .unwrap_or(mlx_sys::mlx_array_new()),
185            group_size,
186            bits,
187            DEFAULT_MODE.as_ptr(),
188            global_scale.as_ptr(),
189            optional_dtype_none(),
190            stream.as_ref().as_ptr(),
191        )
192    })
193}
194
195/// Compatibility shim for [`dequantize`].
196#[generate_macro(customize(forwarding_shim = true))]
197#[deprecated(
198    since = "0.26.0",
199    note = "use `with_stream` or `with_device` around `dequantize`"
200)]
201pub fn dequantize_device<'a>(
202    w: impl AsRef<Array>,
203    scales: impl AsRef<Array>,
204    #[optional] biases: impl Into<Option<&'a Array>>,
205    #[optional] group_size: impl Into<Option<i32>>,
206    #[optional] bits: impl Into<Option<i32>>,
207    #[optional] stream: impl AsRef<Stream>,
208) -> Result<Array> {
209    crate::with_stream(stream.as_ref(), || {
210        dequantize(w, scales, biases, group_size, bits)
211    })
212}
213
214/// Perform quantized matrix multiplication with gathered indices.
215///
216/// This combines the functionality of `gather_mm` and `quantized_matmul`, allowing
217/// matrix multiplication with quantized weights and index gathering along batch dimensions.
218///
219/// # Params
220///
221/// - `x`: Input array
222/// - `w`: Quantized weight matrix
223/// - `scales`: Quantization scales
224/// - `biases`: Optional quantization biases (required for affine mode)
225/// - `lhs_indices`: Optional indices to gather from `x`'s batch dimensions
226/// - `rhs_indices`: Optional indices to gather from `w`'s batch dimensions
227/// - `transpose`: If true, transpose the weight matrix (default: true)
228/// - `group_size`: The quantization group size (default: 64)
229/// - `bits`: The number of bits per element (default: 4)
230/// - `sorted_indices`: If true, indicates the indices are sorted (default: false)
231#[allow(clippy::too_many_arguments)]
232pub fn gather_qmm<'b, 'lhs, 'rhs>(
233    x: impl AsRef<Array>,
234    w: impl AsRef<Array>,
235    scales: impl AsRef<Array>,
236    biases: impl Into<Option<&'b Array>>,
237    lhs_indices: impl Into<Option<&'lhs Array>>,
238    rhs_indices: impl Into<Option<&'rhs Array>>,
239    transpose: impl Into<Option<bool>>,
240    group_size: impl Into<Option<i32>>,
241    bits: impl Into<Option<i32>>,
242    sorted_indices: impl Into<Option<bool>>,
243) -> Result<Array> {
244    let stream = Stream::thread_local_or_default();
245    let transpose = transpose.into().unwrap_or(true);
246    let group_size = optional_int(group_size.into(), DEFAULT_GROUP_SIZE);
247    let bits = optional_int(bits.into(), DEFAULT_BITS);
248    let sorted = sorted_indices.into().unwrap_or(false);
249
250    unsafe {
251        let biases_ptr = biases
252            .into()
253            .map(|a| a.as_ptr())
254            .unwrap_or(mlx_sys::mlx_array_new());
255        let lhs_ptr = lhs_indices
256            .into()
257            .map(|i| i.as_ptr())
258            .unwrap_or(mlx_sys::mlx_array_new());
259        let rhs_ptr = rhs_indices
260            .into()
261            .map(|i| i.as_ptr())
262            .unwrap_or(mlx_sys::mlx_array_new());
263
264        <Array as Guarded>::try_from_op(|res| {
265            mlx_sys::mlx_gather_qmm(
266                res,
267                x.as_ref().as_ptr(),
268                w.as_ref().as_ptr(),
269                scales.as_ref().as_ptr(),
270                biases_ptr,
271                lhs_ptr,
272                rhs_ptr,
273                transpose,
274                group_size,
275                bits,
276                DEFAULT_MODE.as_ptr(),
277                sorted,
278                stream.as_ref().as_ptr(),
279            )
280        })
281    }
282}
283
284/// Compatibility shim for [`gather_qmm`].
285#[allow(clippy::too_many_arguments)]
286#[generate_macro(customize(forwarding_shim = true))]
287#[deprecated(
288    since = "0.26.0",
289    note = "use `with_stream` or `with_device` around `gather_qmm`"
290)]
291pub fn gather_qmm_device<'b, 'lhs, 'rhs>(
292    x: impl AsRef<Array>,
293    w: impl AsRef<Array>,
294    scales: impl AsRef<Array>,
295    #[optional] biases: impl Into<Option<&'b Array>>,
296    #[optional] lhs_indices: impl Into<Option<&'lhs Array>>,
297    #[optional] rhs_indices: impl Into<Option<&'rhs Array>>,
298    #[optional] transpose: impl Into<Option<bool>>,
299    #[optional] group_size: impl Into<Option<i32>>,
300    #[optional] bits: impl Into<Option<i32>>,
301    #[optional] sorted_indices: impl Into<Option<bool>>,
302    #[optional] stream: impl AsRef<Stream>,
303) -> Result<Array> {
304    crate::with_stream(stream.as_ref(), || {
305        gather_qmm(
306            x,
307            w,
308            scales,
309            biases,
310            lhs_indices,
311            rhs_indices,
312            transpose,
313            group_size,
314            bits,
315            sorted_indices,
316        )
317    })
318}
319
320/// Quantized matrix multiplication with quantization of both inputs.
321///
322/// Performs matrix multiplication where `x` is dynamically quantized and `w` is pre-quantized.
323/// This function supports `nvfp4` and `mxfp8` quantization modes.
324///
325/// Note: This function is only supported on GPU with the CUDA backend (Linux with NVIDIA GPU).
326/// It is not available on macOS.
327///
328/// # Params
329///
330/// - `x`: Input matrix to be dynamically quantized
331/// - `w`: Pre-quantized weight matrix
332/// - `w_scales`: Optional scales for the quantized weights (required if `w` is already quantized)
333/// - `group_size`: The quantization group size (default depends on mode: 16 for nvfp4, 32 for mxfp8)
334/// - `bits`: The number of bits per element (default depends on mode: 4 for nvfp4, 8 for mxfp8)
335/// - `mode`: Quantization mode - either "nvfp4" or "mxfp8" (default: "nvfp4")
336#[cfg(not(target_os = "macos"))]
337#[allow(clippy::too_many_arguments)]
338pub fn qqmm<'a>(
339    x: impl AsRef<Array>,
340    w: impl AsRef<Array>,
341    w_scales: impl Into<Option<&'a Array>>,
342    group_size: impl Into<Option<i32>>,
343    bits: impl Into<Option<i32>>,
344    mode: impl Into<Option<&'a str>>,
345) -> Result<Array> {
346    let stream = Stream::thread_local_or_default();
347    let mode_str = mode.into().unwrap_or("nvfp4");
348    let mode_cstr = std::ffi::CString::new(mode_str).expect("Invalid mode string");
349
350    // Defaults depend on mode
351    let (default_group_size, default_bits) = match mode_str {
352        "nvfp4" => (16, 4),
353        "mxfp8" => (32, 8),
354        _ => (16, 4), // fallback to nvfp4 defaults
355    };
356
357    let group_size = optional_int(group_size.into(), default_group_size);
358    let bits = optional_int(bits.into(), default_bits);
359    let global_scale_x = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
360    let global_scale_w = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
361
362    <Array as Guarded>::try_from_op(|res| unsafe {
363        mlx_sys::mlx_qqmm(
364            res,
365            x.as_ref().as_ptr(),
366            w.as_ref().as_ptr(),
367            w_scales
368                .into()
369                .map(|a| a.as_ptr())
370                .unwrap_or(mlx_sys::mlx_array_new()),
371            group_size,
372            bits,
373            mode_cstr.as_ptr(),
374            global_scale_x.as_ptr(),
375            global_scale_w.as_ptr(),
376            stream.as_ref().as_ptr(),
377        )
378    })
379}
380
381/// Compatibility shim for [`qqmm`].
382#[cfg(not(target_os = "macos"))]
383#[allow(clippy::too_many_arguments)]
384#[generate_macro(customize(forwarding_shim = true))]
385#[deprecated(
386    since = "0.26.0",
387    note = "use `with_stream` or `with_device` around `qqmm`"
388)]
389pub fn qqmm_device<'a>(
390    x: impl AsRef<Array>,
391    w: impl AsRef<Array>,
392    #[optional] w_scales: impl Into<Option<&'a Array>>,
393    #[optional] group_size: impl Into<Option<i32>>,
394    #[optional] bits: impl Into<Option<i32>>,
395    #[optional] mode: impl Into<Option<&'a str>>,
396    #[optional] stream: impl AsRef<Stream>,
397) -> Result<Array> {
398    crate::with_stream(stream.as_ref(), || {
399        qqmm(x, w, w_scales, group_size, bits, mode)
400    })
401}
402
403#[cfg(test)]
404mod tests {
405    use crate::{
406        ops::{dequantize, expand_dims, quantize, quantized_matmul},
407        random, Array,
408    };
409
410    #[test]
411    fn test_quantize_dequantize() {
412        let x1 = Array::ones::<f32>(&[128, 1]).unwrap();
413        let x2 = expand_dims(Array::arange::<_, f32>(0, 512, None).unwrap(), 0).unwrap();
414        let x = x1 * x2;
415
416        for i in [2, 4, 8].iter() {
417            let el_per_int = 32 / i;
418            let (x_q, scales, biases) = quantize(&x, 128, *i).unwrap();
419            assert_eq!(x_q.shape(), [128, 512 / el_per_int]);
420            assert_eq!(scales.shape(), [128, 4]);
421            assert_eq!(biases.shape(), [128, 4]);
422
423            let x_hat = dequantize(&x_q, &scales, &biases, 128, *i).unwrap();
424            let max_diff = ((&x - &x_hat).abs().unwrap().max(None).unwrap()).item_exact::<f32>();
425            assert!(max_diff <= 127.0 / (1 << i) as f32);
426        }
427    }
428
429    // Test adapted from Python test `test_quantized.py/test_qmm`
430    #[test]
431    fn test_quantized_matmul() {
432        random::seed(0).unwrap();
433
434        let group_size = 64;
435        let bits = 4;
436        let m = 32;
437        let n = 128;
438        let k = 128;
439
440        let scale = 1.0 / (k as f32).sqrt();
441        let x = random::normal::<f32>(&[m, k], None, None, None).unwrap() * scale;
442        let w = random::normal::<f32>(&[k, n], None, None, None).unwrap() * scale;
443
444        let (w_q, scales, biases) = quantize(&w, group_size, bits).unwrap();
445        let w_hat = dequantize(&w_q, &scales, &biases, group_size, bits).unwrap();
446
447        // Test with biases
448        let y_q = quantized_matmul(&x, &w_q, &scales, &biases, false, group_size, bits).unwrap();
449        let y_hat = x.matmul(&w_hat).unwrap();
450
451        assert_eq!(y_q.shape(), y_hat.shape());
452        let max_diff = ((&y_q - &y_hat).abs().unwrap().max(None).unwrap()).item_exact::<f32>();
453        assert!(max_diff < 1e-3, "max_diff: {}", max_diff);
454    }
455
456    // Test adapted from Python test `test_quantized.py/test_gather_qmm`
457    #[test]
458    fn test_gather_qmm() {
459        use crate::ops::{gather_mm, gather_qmm, swap_axes};
460
461        random::seed(0).unwrap();
462
463        let group_size = 64;
464        let bits = 4;
465
466        // Helper to quantize with transpose option
467        fn quantize_with_transpose(
468            w: &Array,
469            transpose: bool,
470            group_size: i32,
471            bits: i32,
472        ) -> (Array, Array, Array, Array) {
473            let (w_q, scales, biases) = quantize(w, group_size, bits).unwrap();
474            let mut w_hat = dequantize(&w_q, &scales, &biases, group_size, bits).unwrap();
475            if transpose {
476                w_hat = swap_axes(&w_hat, -1, -2).unwrap();
477            }
478            (w_hat, w_q, scales, biases)
479        }
480
481        // Test case 1: batch_A=(1,), lhs_indices=(0,), batch_B=(3,), rhs_indices=(2, 1)
482        let m = 32;
483        let n = 64;
484        let k = 64;
485
486        let x = random::normal::<f32>(&[1, m, k], None, None, None).unwrap();
487        let w = random::normal::<f32>(&[3, n, k], None, None, None).unwrap(); // transpose=true shape
488        let (w_hat, w_q, scales, biases) = quantize_with_transpose(&w, true, group_size, bits);
489
490        let lhs_indices = Array::from_slice(&[0u32], &[1]);
491        let rhs_indices = Array::from_slice(&[2u32, 1], &[2]);
492
493        // Compare gather_mm on dequantized weights vs gather_qmm
494        let c1 = gather_mm(&x, &w_hat, &lhs_indices, &rhs_indices, None).unwrap();
495        let c2 = gather_qmm(
496            &x,
497            &w_q,
498            &scales,
499            &biases,
500            &lhs_indices,
501            &rhs_indices,
502            true,
503            group_size,
504            bits,
505            None,
506        )
507        .unwrap();
508        assert!(
509            c1.all_close(&c2, 1e-4, 1e-4, None).unwrap(),
510            "gather_qmm test case 1 failed"
511        );
512
513        // Test case 2: batch_A=(5,), lhs_indices=(0, 2), batch_B=(3,), rhs_indices=(2, 1)
514        let x = random::normal::<f32>(&[5, m, k], None, None, None).unwrap();
515        let lhs_indices = Array::from_slice(&[0u32, 2], &[2]);
516
517        let c1 = gather_mm(&x, &w_hat, &lhs_indices, &rhs_indices, None).unwrap();
518        let c2 = gather_qmm(
519            &x,
520            &w_q,
521            &scales,
522            &biases,
523            &lhs_indices,
524            &rhs_indices,
525            true,
526            group_size,
527            bits,
528            None,
529        )
530        .unwrap();
531        assert!(
532            c1.all_close(&c2, 1e-4, 1e-4, None).unwrap(),
533            "gather_qmm test case 2 failed"
534        );
535    }
536}
mlx_rs/ops/quantization.rs

mlx_rs/ops/
quantization.rs
