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
11/// Packed weight format. All variants support CPU and Metal in MLX 0.32.2.
12///
13/// Defaults are (group size, bits): affine (64, 4), mxfp4 (32, 4),
14/// mxfp8 (32, 8), and nvfp4 (16, 4). Global scaling is not exposed.
15#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
16pub enum QuantizationMode {
17    /// Integer quantization with a scale and bias per group.
18    #[default]
19    Affine,
20    /// E2M1 values with E8M0 scales.
21    Mxfp4,
22    /// E4M3 values with E8M0 scales.
23    Mxfp8,
24    /// E2M1 values with E4M3 scales, without a global scale.
25    Nvfp4,
26}
27
28/// Invalid quantization configuration or an upstream operation failure.
29#[derive(Debug, thiserror::Error)]
30pub enum QuantizationError {
31    /// The mode does not support this group size and bit width.
32    #[error("unsupported {mode:?} quantization: group size {group_size}, bits {bits}")]
33    InvalidConfiguration {
34        /// Requested format.
35        mode: QuantizationMode,
36        /// Resolved group size.
37        group_size: i32,
38        /// Resolved bit width.
39        bits: i32,
40    },
41    /// Diagnostic from MLX.
42    #[error(transparent)]
43    Exception(#[from] crate::error::Exception),
44}
45
46impl From<QuantizationError> for crate::error::Exception {
47    fn from(error: QuantizationError) -> Self {
48        match error {
49            QuantizationError::Exception(error) => error,
50            error => Self::custom(error.to_string()),
51        }
52    }
53}
54
55impl QuantizationMode {
56    fn validate(
57        self,
58        group_size: Option<i32>,
59        bits: Option<i32>,
60    ) -> std::result::Result<(), QuantizationError> {
61        let (default_group, default_bits) = match self {
62            Self::Affine => (64, 4),
63            Self::Mxfp4 => (32, 4),
64            Self::Mxfp8 => (32, 8),
65            Self::Nvfp4 => (16, 4),
66        };
67        let group_size = group_size.unwrap_or(default_group);
68        let bits = bits.unwrap_or(default_bits);
69        let valid = match self {
70            Self::Affine => {
71                matches!(group_size, 32 | 64 | 128) && matches!(bits, 2 | 3 | 4 | 5 | 6 | 8)
72            }
73            _ => group_size == default_group && bits == default_bits,
74        };
75        if !valid {
76            return Err(QuantizationError::InvalidConfiguration {
77                mode: self,
78                group_size,
79                bits,
80            });
81        }
82        Ok(())
83    }
84
85    fn as_c_str(self) -> &'static CStr {
86        match self {
87            Self::Affine => c"affine",
88            Self::Mxfp4 => c"mxfp4",
89            Self::Mxfp8 => c"mxfp8",
90            Self::Nvfp4 => c"nvfp4",
91        }
92    }
93}
94
95/// Configuration for [`quantize`]. `None` selects the mode's default.
96#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
97pub struct QuantizeOptions {
98    /// Quantization format.
99    pub mode: QuantizationMode,
100    /// Elements sharing a scale: affine 32/64/128, MX 32, NV 16.
101    pub group_size: Option<i32>,
102    /// Bits per element: affine 2/3/4/5/6/8, mxfp8 8, other modes 4.
103    pub bits: Option<i32>,
104}
105
106/// Configuration for [`dequantize`]. `None` selects the upstream mode default.
107#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
108pub struct DequantizeOptions {
109    /// Quantization format used to pack the weights.
110    pub mode: QuantizationMode,
111    /// Elements sharing a scale; see [`QuantizeOptions`].
112    pub group_size: Option<i32>,
113    /// Bits per packed element; see [`QuantizeOptions`].
114    pub bits: Option<i32>,
115}
116
117/// Configuration for [`quantized_matmul`].
118#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
119pub struct QuantizedMatmulOptions {
120    /// Quantization format used to pack the weights.
121    pub mode: QuantizationMode,
122    /// Elements sharing a scale; `None` selects the mode's default.
123    pub group_size: Option<i32>,
124    /// Bits per packed element; `None` selects the mode's default.
125    pub bits: Option<i32>,
126    /// Transpose the packed weight matrix. Defaults to false.
127    pub transpose: bool,
128}
129
130/// Configuration for [`gather_qmm`].
131#[derive(Debug, Clone, Copy)]
132pub struct GatherQmmOptions<'a> {
133    /// Quantization format used to pack the weights.
134    pub mode: QuantizationMode,
135    /// Elements sharing a scale; `None` selects the mode's default.
136    pub group_size: Option<i32>,
137    /// Bits per packed element; `None` selects the mode's default.
138    pub bits: Option<i32>,
139    /// Transpose the packed weight matrix. Defaults to true.
140    pub transpose: bool,
141    /// Indices into the input batch dimensions.
142    pub lhs_indices: Option<&'a Array>,
143    /// Indices into the weight batch dimensions.
144    pub rhs_indices: Option<&'a Array>,
145    /// Whether the supplied indices are sorted. Defaults to false.
146    pub sorted_indices: bool,
147}
148
149impl Default for GatherQmmOptions<'_> {
150    fn default() -> Self {
151        Self {
152            mode: QuantizationMode::default(),
153            group_size: None,
154            bits: None,
155            transpose: true,
156            lhs_indices: None,
157            rhs_indices: None,
158            sorted_indices: false,
159        }
160    }
161}
162
163const DEFAULT_GROUP_SIZE: i32 = 64;
164const DEFAULT_BITS: i32 = 4;
165
166/// Helper to convert Option<i32> to mlx_optional_int
167fn optional_int(value: Option<i32>, default: i32) -> mlx_sys::mlx_optional_int {
168    mlx_sys::mlx_optional_int {
169        value: value.unwrap_or(default),
170        has_value: value.is_some(),
171    }
172}
173
174/// Helper to create a "no value" optional dtype
175fn optional_dtype_none() -> mlx_sys::mlx_optional_dtype {
176    mlx_sys::mlx_optional_dtype {
177        value: mlx_sys::mlx_dtype__MLX_FLOAT32, // default value, ignored when has_value is false
178        has_value: false,
179    }
180}
181
182/// Quantize weights into packed values, scales, and optional biases.
183///
184/// Biases are present only for [`QuantizationMode::Affine`]. Inputs must have at
185/// least two dimensions and a last dimension divisible by the group size.
186/// Invalid mode/group/bit combinations return [`QuantizationError::InvalidConfiguration`].
187/// Shapes and dtypes are checked upstream.
188///
189/// ```
190/// use mlx_rs::{Array, ops::{self, QuantizationMode, QuantizeOptions, DequantizeOptions}};
191/// # fn main() -> Result<(), ops::QuantizationError> {
192/// let weights = Array::ones::<f32>(&[32, 64])?;
193/// let mode = QuantizationMode::Mxfp4;
194/// let (packed, scales, biases) = ops::quantize(
195///     &weights, QuantizeOptions { mode, ..Default::default() },
196/// )?;
197/// assert!(biases.is_none());
198/// let restored = ops::dequantize(
199///     &packed, &scales, biases.as_ref(),
200///     DequantizeOptions { mode, ..Default::default() },
201/// )?;
202/// assert_eq!(restored.shape(), weights.shape());
203/// # Ok(())
204/// # }
205/// ```
206pub fn quantize(
207    w: impl AsRef<Array>,
208    options: QuantizeOptions,
209) -> std::result::Result<(Array, Array, Option<Array>), QuantizationError> {
210    options.mode.validate(options.group_size, options.bits)?;
211    let stream = Stream::thread_local_or_default();
212    let group_size = optional_int(options.group_size, DEFAULT_GROUP_SIZE);
213    let bits = optional_int(options.bits, DEFAULT_BITS);
214    let global_scale = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
215
216    let result = VectorArray::try_from_op(|res| unsafe {
217        mlx_sys::mlx_quantize(
218            res,
219            w.as_ref().as_ptr(),
220            group_size,
221            bits,
222            options.mode.as_c_str().as_ptr(),
223            global_scale.as_ptr(),
224            stream.as_ref().as_ptr(),
225        )
226    })?;
227
228    let arrays: Vec<Array> = result.try_into_values()?;
229    let expected = if options.mode == QuantizationMode::Affine {
230        3
231    } else {
232        2
233    };
234    if arrays.len() != expected {
235        return Err(crate::error::Exception::custom(format!(
236            "Expected {expected} arrays from quantize, got {}",
237            arrays.len()
238        ))
239        .into());
240    }
241    let mut iter = arrays.into_iter();
242    Ok((iter.next().unwrap(), iter.next().unwrap(), iter.next()))
243}
244
245/// Compatibility shim for [`quantize`].
246#[generate_macro(customize(forwarding_shim = true))]
247#[deprecated(
248    since = "0.26.0",
249    note = "use `with_stream` or `with_device` around `quantize`"
250)]
251pub fn quantize_device(
252    w: impl AsRef<Array>,
253    #[optional] group_size: impl Into<Option<i32>>,
254    #[optional] bits: impl Into<Option<i32>>,
255    #[optional] stream: impl AsRef<Stream>,
256) -> Result<(Array, Array, Array)> {
257    crate::with_stream(stream.as_ref(), || {
258        let (weights, scales, biases) = quantize(
259            w,
260            QuantizeOptions {
261                group_size: group_size.into(),
262                bits: bits.into(),
263                ..Default::default()
264            },
265        )?;
266        Ok((
267            weights,
268            scales,
269            biases.expect("affine quantization returns biases"),
270        ))
271    })
272}
273
274/// Multiply `x` by packed weights. Pass biases for affine weights and `None`
275/// for floating-point modes; options must match the weights' quantization.
276pub fn quantized_matmul<'a>(
277    x: impl AsRef<Array>,
278    w: impl AsRef<Array>,
279    scales: impl AsRef<Array>,
280    biases: impl Into<Option<&'a Array>>,
281    options: QuantizedMatmulOptions,
282) -> std::result::Result<Array, QuantizationError> {
283    options.mode.validate(options.group_size, options.bits)?;
284    let stream = Stream::thread_local_or_default();
285    let absent = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
286    let transpose = options.transpose;
287    let group_size = optional_int(options.group_size, DEFAULT_GROUP_SIZE);
288    let bits = optional_int(options.bits, DEFAULT_BITS);
289
290    <Array as Guarded>::try_from_op(|res| unsafe {
291        mlx_sys::mlx_quantized_matmul(
292            res,
293            x.as_ref().as_ptr(),
294            w.as_ref().as_ptr(),
295            scales.as_ref().as_ptr(),
296            biases.into().map(|a| a.as_ptr()).unwrap_or(absent.as_ptr()),
297            transpose,
298            group_size,
299            bits,
300            options.mode.as_c_str().as_ptr(),
301            stream.as_ref().as_ptr(),
302        )
303    })
304    .map_err(Into::into)
305}
306
307/// Compatibility shim for [`quantized_matmul`].
308#[allow(clippy::too_many_arguments)]
309#[generate_macro(customize(forwarding_shim = true))]
310#[deprecated(
311    since = "0.26.0",
312    note = "use `with_stream` or `with_device` around `quantized_matmul`"
313)]
314pub fn quantized_matmul_device<'a>(
315    x: impl AsRef<Array>,
316    w: impl AsRef<Array>,
317    scales: impl AsRef<Array>,
318    #[optional] biases: impl Into<Option<&'a Array>>,
319    #[optional] transpose: impl Into<Option<bool>>,
320    #[optional] group_size: impl Into<Option<i32>>,
321    #[optional] bits: impl Into<Option<i32>>,
322    #[optional] stream: impl AsRef<Stream>,
323) -> Result<Array> {
324    crate::with_stream(stream.as_ref(), || {
325        quantized_matmul(
326            x,
327            w,
328            scales,
329            biases,
330            QuantizedMatmulOptions {
331                transpose: transpose.into().unwrap_or(false),
332                group_size: group_size.into(),
333                bits: bits.into(),
334                ..Default::default()
335            },
336        )
337        .map_err(Into::into)
338    })
339}
340
341/// Reconstruct weights using the same mode, group size, and bits used to quantize them.
342/// Pass biases for affine weights and `None` for the floating-point modes.
343/// Floating-point modes dequantize to bfloat16; affine uses the scale/bias dtype.
344pub fn dequantize<'a>(
345    w: impl AsRef<Array>,
346    scales: impl AsRef<Array>,
347    biases: impl Into<Option<&'a Array>>,
348    options: DequantizeOptions,
349) -> std::result::Result<Array, QuantizationError> {
350    options.mode.validate(options.group_size, options.bits)?;
351    let stream = Stream::thread_local_or_default();
352    let absent = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
353    let group_size = optional_int(options.group_size, DEFAULT_GROUP_SIZE);
354    let bits = optional_int(options.bits, DEFAULT_BITS);
355    let global_scale = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
356
357    <Array as Guarded>::try_from_op(|res| unsafe {
358        mlx_sys::mlx_dequantize(
359            res,
360            w.as_ref().as_ptr(),
361            scales.as_ref().as_ptr(),
362            biases.into().map(|a| a.as_ptr()).unwrap_or(absent.as_ptr()),
363            group_size,
364            bits,
365            options.mode.as_c_str().as_ptr(),
366            global_scale.as_ptr(),
367            optional_dtype_none(),
368            stream.as_ref().as_ptr(),
369        )
370    })
371    .map_err(Into::into)
372}
373
374/// Compatibility shim for [`dequantize`].
375#[generate_macro(customize(forwarding_shim = true))]
376#[deprecated(
377    since = "0.26.0",
378    note = "use `with_stream` or `with_device` around `dequantize`"
379)]
380pub fn dequantize_device<'a>(
381    w: impl AsRef<Array>,
382    scales: impl AsRef<Array>,
383    #[optional] biases: impl Into<Option<&'a Array>>,
384    #[optional] group_size: impl Into<Option<i32>>,
385    #[optional] bits: impl Into<Option<i32>>,
386    #[optional] stream: impl AsRef<Stream>,
387) -> Result<Array> {
388    crate::with_stream(stream.as_ref(), || {
389        dequantize(
390            w,
391            scales,
392            biases,
393            DequantizeOptions {
394                group_size: group_size.into(),
395                bits: bits.into(),
396                ..Default::default()
397            },
398        )
399        .map_err(Into::into)
400    })
401}
402
403/// Multiply selected input and packed-weight batches.
404/// Pass affine biases or `None` for floating-point modes. Quantization options
405/// must match the packed weights; indices select the respective batch dimensions.
406pub fn gather_qmm<'b>(
407    x: impl AsRef<Array>,
408    w: impl AsRef<Array>,
409    scales: impl AsRef<Array>,
410    biases: impl Into<Option<&'b Array>>,
411    options: GatherQmmOptions<'_>,
412) -> std::result::Result<Array, QuantizationError> {
413    options.mode.validate(options.group_size, options.bits)?;
414    let stream = Stream::thread_local_or_default();
415    let absent = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
416    let transpose = options.transpose;
417    let group_size = optional_int(options.group_size, DEFAULT_GROUP_SIZE);
418    let bits = optional_int(options.bits, DEFAULT_BITS);
419    let sorted = options.sorted_indices;
420
421    unsafe {
422        let biases_ptr = biases.into().map(|a| a.as_ptr()).unwrap_or(absent.as_ptr());
423        let lhs_ptr = options
424            .lhs_indices
425            .map(|i| i.as_ptr())
426            .unwrap_or(absent.as_ptr());
427        let rhs_ptr = options
428            .rhs_indices
429            .map(|i| i.as_ptr())
430            .unwrap_or(absent.as_ptr());
431
432        <Array as Guarded>::try_from_op(|res| {
433            mlx_sys::mlx_gather_qmm(
434                res,
435                x.as_ref().as_ptr(),
436                w.as_ref().as_ptr(),
437                scales.as_ref().as_ptr(),
438                biases_ptr,
439                lhs_ptr,
440                rhs_ptr,
441                transpose,
442                group_size,
443                bits,
444                options.mode.as_c_str().as_ptr(),
445                sorted,
446                stream.as_ref().as_ptr(),
447            )
448        })
449        .map_err(Into::into)
450    }
451}
452
453/// Compatibility shim for [`gather_qmm`].
454#[allow(clippy::too_many_arguments)]
455#[generate_macro(customize(forwarding_shim = true))]
456#[deprecated(
457    since = "0.26.0",
458    note = "use `with_stream` or `with_device` around `gather_qmm`"
459)]
460pub fn gather_qmm_device<'b, 'lhs, 'rhs>(
461    x: impl AsRef<Array>,
462    w: impl AsRef<Array>,
463    scales: impl AsRef<Array>,
464    #[optional] biases: impl Into<Option<&'b Array>>,
465    #[optional] lhs_indices: impl Into<Option<&'lhs Array>>,
466    #[optional] rhs_indices: impl Into<Option<&'rhs Array>>,
467    #[optional] transpose: impl Into<Option<bool>>,
468    #[optional] group_size: impl Into<Option<i32>>,
469    #[optional] bits: impl Into<Option<i32>>,
470    #[optional] sorted_indices: impl Into<Option<bool>>,
471    #[optional] stream: impl AsRef<Stream>,
472) -> Result<Array> {
473    crate::with_stream(stream.as_ref(), || {
474        gather_qmm(
475            x,
476            w,
477            scales,
478            biases,
479            GatherQmmOptions {
480                lhs_indices: lhs_indices.into(),
481                rhs_indices: rhs_indices.into(),
482                transpose: transpose.into().unwrap_or(true),
483                group_size: group_size.into(),
484                bits: bits.into(),
485                sorted_indices: sorted_indices.into().unwrap_or(false),
486                ..Default::default()
487            },
488        )
489        .map_err(Into::into)
490    })
491}
492
493/// Quantized matrix multiplication with quantization of both inputs.
494///
495/// Performs matrix multiplication where `x` is dynamically quantized and `w` is pre-quantized.
496/// This function supports `nvfp4` and `mxfp8` quantization modes.
497///
498/// Note: This function is only supported on GPU with the CUDA backend (Linux with NVIDIA GPU).
499/// It is not available on macOS.
500///
501/// # Params
502///
503/// - `x`: Input matrix to be dynamically quantized
504/// - `w`: Pre-quantized weight matrix
505/// - `w_scales`: Optional scales for the quantized weights (required if `w` is already quantized)
506/// - `group_size`: The quantization group size (default depends on mode: 16 for nvfp4, 32 for mxfp8)
507/// - `bits`: The number of bits per element (default depends on mode: 4 for nvfp4, 8 for mxfp8)
508/// - `mode`: Quantization mode - either "nvfp4" or "mxfp8" (default: "nvfp4")
509#[cfg(not(target_os = "macos"))]
510#[allow(clippy::too_many_arguments)]
511pub fn qqmm<'a>(
512    x: impl AsRef<Array>,
513    w: impl AsRef<Array>,
514    w_scales: impl Into<Option<&'a Array>>,
515    group_size: impl Into<Option<i32>>,
516    bits: impl Into<Option<i32>>,
517    mode: impl Into<Option<&'a str>>,
518) -> Result<Array> {
519    let stream = Stream::thread_local_or_default();
520    let mode_str = mode.into().unwrap_or("nvfp4");
521    let mode_cstr = std::ffi::CString::new(mode_str).expect("Invalid mode string");
522
523    // Defaults depend on mode
524    let (default_group_size, default_bits) = match mode_str {
525        "nvfp4" => (16, 4),
526        "mxfp8" => (32, 8),
527        _ => (16, 4), // fallback to nvfp4 defaults
528    };
529
530    let group_size = optional_int(group_size.into(), default_group_size);
531    let bits = optional_int(bits.into(), default_bits);
532    let global_scale_x = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
533    let global_scale_w = unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) };
534
535    <Array as Guarded>::try_from_op(|res| unsafe {
536        mlx_sys::mlx_qqmm(
537            res,
538            x.as_ref().as_ptr(),
539            w.as_ref().as_ptr(),
540            w_scales
541                .into()
542                .map(|a| a.as_ptr())
543                .unwrap_or(mlx_sys::mlx_array_new()),
544            group_size,
545            bits,
546            mode_cstr.as_ptr(),
547            global_scale_x.as_ptr(),
548            global_scale_w.as_ptr(),
549            stream.as_ref().as_ptr(),
550        )
551    })
552}
553
554/// Compatibility shim for [`qqmm`].
555#[cfg(not(target_os = "macos"))]
556#[allow(clippy::too_many_arguments)]
557#[generate_macro(customize(forwarding_shim = true))]
558#[deprecated(
559    since = "0.26.0",
560    note = "use `with_stream` or `with_device` around `qqmm`"
561)]
562pub fn qqmm_device<'a>(
563    x: impl AsRef<Array>,
564    w: impl AsRef<Array>,
565    #[optional] w_scales: impl Into<Option<&'a Array>>,
566    #[optional] group_size: impl Into<Option<i32>>,
567    #[optional] bits: impl Into<Option<i32>>,
568    #[optional] mode: impl Into<Option<&'a str>>,
569    #[optional] stream: impl AsRef<Stream>,
570) -> Result<Array> {
571    crate::with_stream(stream.as_ref(), || {
572        qqmm(x, w, w_scales, group_size, bits, mode)
573    })
574}
575
576#[cfg(test)]
577mod tests {
578    use crate::{
579        ops::{dequantize, expand_dims, quantize, quantized_matmul},
580        random, Array,
581    };
582
583    #[test]
584    fn invalid_quantization_options_do_not_reach_ffi() {
585        use super::{QuantizationError, QuantizationMode, QuantizeOptions, QuantizedMatmulOptions};
586        let weights = Array::ones::<f32>(&[32, 64]).unwrap();
587        for mode in [
588            QuantizationMode::Affine,
589            QuantizationMode::Mxfp4,
590            QuantizationMode::Mxfp8,
591            QuantizationMode::Nvfp4,
592        ] {
593            for (group_size, bits) in [(Some(0), None), (None, Some(0)), (Some(7), Some(7))] {
594                assert!(matches!(
595                    quantize(
596                        &weights,
597                        QuantizeOptions {
598                            mode,
599                            group_size,
600                            bits
601                        }
602                    ),
603                    Err(QuantizationError::InvalidConfiguration { .. })
604                ));
605                assert!(matches!(
606                    quantized_matmul(
607                        &weights,
608                        &weights,
609                        &weights,
610                        None,
611                        QuantizedMatmulOptions {
612                            mode,
613                            group_size,
614                            bits,
615                            ..Default::default()
616                        }
617                    ),
618                    Err(QuantizationError::InvalidConfiguration { .. })
619                ));
620            }
621        }
622    }
623
624    #[test]
625    fn test_quantize_dequantize() {
626        let x1 = Array::ones::<f32>(&[128, 1]).unwrap();
627        let x2 = expand_dims(Array::arange::<_, f32>(0, 512, None).unwrap(), 0).unwrap();
628        let x = x1 * x2;
629
630        for i in [2, 4, 8].iter() {
631            let el_per_int = 32 / i;
632            let (x_q, scales, biases) = quantize(
633                &x,
634                crate::ops::QuantizeOptions {
635                    group_size: Some(128),
636                    bits: Some(*i),
637                    ..Default::default()
638                },
639            )
640            .unwrap();
641            let biases = biases.expect("affine quantization returns biases");
642            assert_eq!(x_q.shape(), [128, 512 / el_per_int]);
643            assert_eq!(scales.shape(), [128, 4]);
644            assert_eq!(biases.shape(), [128, 4]);
645
646            let x_hat = dequantize(
647                &x_q,
648                &scales,
649                &biases,
650                crate::ops::DequantizeOptions {
651                    group_size: Some(128),
652                    bits: Some(*i),
653                    ..Default::default()
654                },
655            )
656            .unwrap();
657            let max_diff = ((&x - &x_hat).abs().unwrap().max(None).unwrap()).item_exact::<f32>();
658            assert!(max_diff <= 127.0 / (1 << i) as f32);
659        }
660    }
661
662    // Test adapted from Python test `test_quantized.py/test_qmm`
663    #[test]
664    fn test_quantized_matmul() {
665        random::seed(0).unwrap();
666
667        let group_size = 64;
668        let bits = 4;
669        let m = 32;
670        let n = 128;
671        let k = 128;
672
673        let scale = 1.0 / (k as f32).sqrt();
674        let x = random::normal::<f32>(&[m, k], None, None, None).unwrap() * scale;
675        let w = random::normal::<f32>(&[k, n], None, None, None).unwrap() * scale;
676
677        let (w_q, scales, biases) = quantize(
678            &w,
679            crate::ops::QuantizeOptions {
680                group_size: Some(group_size),
681                bits: Some(bits),
682                ..Default::default()
683            },
684        )
685        .unwrap();
686        let biases = biases.expect("affine quantization returns biases");
687        let w_hat = dequantize(
688            &w_q,
689            &scales,
690            &biases,
691            crate::ops::DequantizeOptions {
692                group_size: Some(group_size),
693                bits: Some(bits),
694                ..Default::default()
695            },
696        )
697        .unwrap();
698
699        // Test with biases
700        let y_q = quantized_matmul(
701            &x,
702            &w_q,
703            &scales,
704            &biases,
705            crate::ops::QuantizedMatmulOptions {
706                transpose: false,
707                group_size: Some(group_size),
708                bits: Some(bits),
709                ..Default::default()
710            },
711        )
712        .unwrap();
713        let y_hat = x.matmul(&w_hat).unwrap();
714
715        assert_eq!(y_q.shape(), y_hat.shape());
716        let max_diff = ((&y_q - &y_hat).abs().unwrap().max(None).unwrap()).item_exact::<f32>();
717        assert!(max_diff < 1e-3, "max_diff: {}", max_diff);
718    }
719
720    // Test adapted from Python test `test_quantized.py/test_gather_qmm`
721    #[test]
722    fn test_gather_qmm() {
723        use crate::ops::{gather_mm, gather_qmm, swap_axes};
724
725        random::seed(0).unwrap();
726
727        let group_size = 64;
728        let bits = 4;
729
730        // Helper to quantize with transpose option
731        fn quantize_with_transpose(
732            w: &Array,
733            transpose: bool,
734            group_size: i32,
735            bits: i32,
736        ) -> (Array, Array, Array, Array) {
737            let (w_q, scales, biases) = quantize(
738                w,
739                crate::ops::QuantizeOptions {
740                    group_size: Some(group_size),
741                    bits: Some(bits),
742                    ..Default::default()
743                },
744            )
745            .unwrap();
746            let biases = biases.expect("affine quantization returns biases");
747            let mut w_hat = dequantize(
748                &w_q,
749                &scales,
750                &biases,
751                crate::ops::DequantizeOptions {
752                    group_size: Some(group_size),
753                    bits: Some(bits),
754                    ..Default::default()
755                },
756            )
757            .unwrap();
758            if transpose {
759                w_hat = swap_axes(&w_hat, -1, -2).unwrap();
760            }
761            (w_hat, w_q, scales, biases)
762        }
763
764        // Test case 1: batch_A=(1,), lhs_indices=(0,), batch_B=(3,), rhs_indices=(2, 1)
765        let m = 32;
766        let n = 64;
767        let k = 64;
768
769        let x = random::normal::<f32>(&[1, m, k], None, None, None).unwrap();
770        let w = random::normal::<f32>(&[3, n, k], None, None, None).unwrap(); // transpose=true shape
771        let (w_hat, w_q, scales, biases) = quantize_with_transpose(&w, true, group_size, bits);
772
773        let lhs_indices = Array::from_slice(&[0u32], &[1]);
774        let rhs_indices = Array::from_slice(&[2u32, 1], &[2]);
775
776        // Compare gather_mm on dequantized weights vs gather_qmm
777        let c1 = gather_mm(&x, &w_hat, &lhs_indices, &rhs_indices, None).unwrap();
778        let c2 = gather_qmm(
779            &x,
780            &w_q,
781            &scales,
782            &biases,
783            crate::ops::GatherQmmOptions {
784                lhs_indices: Some(&lhs_indices),
785                rhs_indices: Some(&rhs_indices),
786                transpose: true,
787                group_size: Some(group_size),
788                bits: Some(bits),
789                ..Default::default()
790            },
791        )
792        .unwrap();
793        assert!(
794            c1.all_close(&c2, 1e-4, 1e-4, None).unwrap(),
795            "gather_qmm test case 1 failed"
796        );
797
798        // Test case 2: batch_A=(5,), lhs_indices=(0, 2), batch_B=(3,), rhs_indices=(2, 1)
799        let x = random::normal::<f32>(&[5, m, k], None, None, None).unwrap();
800        let lhs_indices = Array::from_slice(&[0u32, 2], &[2]);
801
802        let c1 = gather_mm(&x, &w_hat, &lhs_indices, &rhs_indices, None).unwrap();
803        let c2 = gather_qmm(
804            &x,
805            &w_q,
806            &scales,
807            &biases,
808            crate::ops::GatherQmmOptions {
809                lhs_indices: Some(&lhs_indices),
810                rhs_indices: Some(&rhs_indices),
811                transpose: true,
812                group_size: Some(group_size),
813                bits: Some(bits),
814                ..Default::default()
815            },
816        )
817        .unwrap();
818        assert!(
819            c1.all_close(&c2, 1e-4, 1e-4, None).unwrap(),
820            "gather_qmm test case 2 failed"
821        );
822    }
823}
mlx_rs/ops/quantization.rs

mlx_rs/ops/
quantization.rs
