nuvec_core/
quantized.rs

1//! Affine quantization, as in MLX: weights are stored with `bits` bits per value and one
2//! scale and bias per group of `group_size` values along the last axis.
3//!
4//! The packed form is MLX's, so weights quantized by either library load in the other:
5//! codes form a little-endian bit stream per row, stored as uint32; scales and biases keep
6//! the weights' dtype. A value `w` is stored as `round((w - bias) / scale)`, clamped to
7//! `[0, 2^bits - 1]`.
8//!
9//! ```
10//! use nuvec_core::{Array, Result, dequantize, quantize, quantized_matmul};
11//!
12//! # fn main() -> Result<()> {
13//! let w = Array::from_slice(&(0..256).map(|i| i as f32 / 64.0).collect::<Vec<_>>(), &[4, 64])?;
14//! let (wq, scales, biases) = quantize(&w, 64, 4)?;
15//! assert_eq!(wq.shape(), &[4, 8]); // 64 values · 4 bits = 8 words per row
16//! let restored = dequantize(&wq, &scales, &biases, 64, 4)?;
17//! assert_eq!(restored.shape(), w.shape());
18//!
19//! let x = Array::full(&[2, 64], 1.0f32)?;
20//! let y = quantized_matmul(&x, &wq, &scales, &biases, true, 64, 4)?; // x · Wᵀ
21//! assert_eq!(y.shape(), &[2, 4]);
22//! # Ok(())
23//! # }
24//! ```
25
26use crate::ops::{astype, broadcast_to, expand_dims, squeeze};
27use crate::primitive::Primitive;
28use crate::{Array, Category, Error, Result, promote_types, view};
29
30/// The supported group sizes and bit widths (MLX's).
31fn check_params(op: &str, group_size: usize, bits: u32) -> Result<u8> {
32    if !matches!(group_size, 32 | 64 | 128) {
33        return Err(Error::InvalidArgument(format!(
34            "{op}: group size {group_size} is not supported (32, 64 or 128)"
35        )));
36    }
37    match bits {
38        1..=6 | 8 => Ok(bits as u8),
39        _ => Err(Error::InvalidArgument(format!(
40            "{op}: {bits} bits are not supported (1, 2, 3, 4, 5, 6 or 8)"
41        ))),
42    }
43}
44
45fn check_inputs(inputs: &[&Array]) -> Result<()> {
46    match inputs.iter().find_map(|a| a.error()) {
47        Some(err) => Err(err),
48        None => Ok(()),
49    }
50}
51
52/// Quantizes the last axis of `w` (float16, bfloat16 or float32, at least 2-D, the last
53/// dimension divisible by `group_size`) in groups of `group_size` (32, 64 or 128) values with
54/// `bits` (1–6 or 8) bits each. Returns the packed codes (uint32, last dimension
55/// `· bits / 32`), the scales and the biases (`w`'s dtype, last dimension `/ group_size`),
56/// equal to MLX's bit for bit.
57pub fn quantize(w: &Array, group_size: usize, bits: u32) -> Result<(Array, Array, Array)> {
58    check_inputs(&[w])?;
59    let bits = check_params("quantize", group_size, bits)?;
60    let mut out = Array::from_op(Primitive::Quantize { group_size, bits }, vec![w.clone()])?;
61    let biases = out.pop().expect("three outputs");
62    let scales = out.pop().expect("three outputs");
63    Ok((out.pop().expect("three outputs"), scales, biases))
64}
65
66/// Restores the values packed by [`quantize`]: `code · scale + bias`, in the scales' dtype.
67pub fn dequantize(
68    w: &Array,
69    scales: &Array,
70    biases: &Array,
71    group_size: usize,
72    bits: u32,
73) -> Result<Array> {
74    check_inputs(&[w, scales, biases])?;
75    let bits = check_params("dequantize", group_size, bits)?;
76    let inputs = vec![w.clone(), scales.clone(), biases.clone()];
77    Ok(Array::from_op(Primitive::Dequantize { group_size, bits }, inputs)?.swap_remove(0))
78}
79
80/// The matrix product of `x` with quantized weights, `x · dequantize(w)ᵀ` if `transpose`
81/// (MLX's default, for weights stored as `[out, in]`) or `x · dequantize(w)`, without
82/// materialising the dequantized matrix.
83///
84/// As in MLX, `x`, `scales` and `biases` are promoted to one floating-point type, and batch
85/// dimensions broadcast when `x` and `w` both have more than two. A one-dimensional `x` is a
86/// row vector.
87#[allow(clippy::too_many_arguments)]
88pub fn quantized_matmul(
89    x: &Array,
90    w: &Array,
91    scales: &Array,
92    biases: &Array,
93    transpose: bool,
94    group_size: usize,
95    bits: u32,
96) -> Result<Array> {
97    check_inputs(&[x, w, scales, biases])?;
98    let bits = check_params("quantized_matmul", group_size, bits)?;
99    let dtype = promote_types(promote_types(x.dtype(), scales.dtype()), biases.dtype());
100    if dtype.category() != Category::Float || dtype.size() > 4 {
101        return Err(Error::UnsupportedDType {
102            op: "quantized_matmul",
103            dtype,
104        });
105    }
106    let (x, scales, biases) = (
107        astype(x, dtype)?,
108        astype(scales, dtype)?,
109        astype(biases, dtype)?,
110    );
111    let vector = x.ndim() == 1;
112    let mut x = if vector { expand_dims(&x, 0)? } else { x };
113    let (mut w, mut scales, mut biases) = (w.clone(), scales, biases);
114    let (xn, wn) = (x.ndim(), w.ndim());
115    if wn > 2 && xn >= 2 {
116        // Broadcast the batch dimensions of every operand together (MLX broadcasts them only
117        // when both have batch dimensions; nuvec also lifts a 2-D `x` to `w`'s batch).
118        let batch = view::broadcast_shapes(
119            "quantized_matmul",
120            &x.shape()[..xn - 2],
121            &w.shape()[..wn - 2],
122        )?;
123        let full = |a: &Array| -> Result<Array> {
124            let mut shape = batch.to_vec();
125            shape.extend_from_slice(&a.shape()[a.ndim() - 2..]);
126            broadcast_to(a, &shape)
127        };
128        x = full(&x)?;
129        w = full(&w)?;
130        scales = full(&scales)?;
131        biases = full(&biases)?;
132    }
133    let primitive = Primitive::QuantizedMatmul {
134        group_size,
135        bits,
136        transpose,
137    };
138    let out = Array::from_op(primitive, vec![x, w, scales, biases])?.swap_remove(0);
139    if vector {
140        squeeze(&out, &[-2])
141    } else {
142        Ok(out)
143    }
144}
145
146/// Quantized matrix products of gathered batches, as MLX's `gather_qmm` (mixture-of-experts
147/// layers): output batch `i` is `x[lhs_indices[i]] · dequantize(w)[rhs_indices[i]]ᵀ` (or
148/// without the transpose), indexing the flattened batch dimensions of `x` (`[..., M, K]`)
149/// and of `w`, `scales` and `biases`. The indices (integers) broadcast together; either
150/// defaults to every batch of its operand in order. `sorted_indices` promises that
151/// `rhs_indices` is sorted (as after grouping rows by expert), which lets one kernel walk
152/// each weight batch once.
153///
154/// ```
155/// use nuvec_core::{Array, Result, gather_qmm, quantize};
156///
157/// # fn main() -> Result<()> {
158/// // Two "experts" of shape [4, 64]: expert e has every weight equal to e + 1.
159/// let w = Array::from_slice(
160///     &(0..512).map(|i| if i < 256 { 1.0f32 } else { 2.0 }).collect::<Vec<_>>(),
161///     &[2, 4, 64],
162/// )?;
163/// let (wq, scales, biases) = quantize(&w, 64, 4)?;
164/// let x = Array::full(&[3, 1, 64], 1.0f32)?; // three tokens
165/// let experts = Array::from_slice(&[1u32, 0, 1], &[3])?;
166/// let y = gather_qmm(&x, &wq, &scales, &biases, None, Some(&experts), true, 64, 4, false)?;
167/// assert_eq!(y.shape(), &[3, 1, 4]);
168/// assert_eq!(y.to_vec::<f32>()?[..5], [128.0, 128.0, 128.0, 128.0, 64.0]);
169/// # Ok(())
170/// # }
171/// ```
172#[allow(clippy::too_many_arguments)]
173pub fn gather_qmm(
174    x: &Array,
175    w: &Array,
176    scales: &Array,
177    biases: &Array,
178    lhs_indices: Option<&Array>,
179    rhs_indices: Option<&Array>,
180    transpose: bool,
181    group_size: usize,
182    bits: u32,
183    sorted_indices: bool,
184) -> Result<Array> {
185    check_inputs(&[x, w, scales, biases])?;
186    let bits = check_params("gather_qmm", group_size, bits)?;
187    let invalid = |msg: String| Err(Error::InvalidArgument(format!("gather_qmm: {msg}")));
188    let dtype = promote_types(promote_types(x.dtype(), scales.dtype()), biases.dtype());
189    if dtype.category() != Category::Float || dtype.size() > 4 {
190        return Err(Error::UnsupportedDType {
191            op: "gather_qmm",
192            dtype,
193        });
194    }
195    if x.ndim() < 2 || w.ndim() < 2 || scales.shape() != biases.shape() {
196        return invalid(format!(
197            "x must be [..., M, K] and w [..., N, K'], got {:?} and {:?}",
198            x.shape(),
199            w.shape()
200        ));
201    }
202    let (x, scales, biases) = (
203        astype(x, dtype)?,
204        astype(scales, dtype)?,
205        astype(biases, dtype)?,
206    );
207    let (lhs, rhs) = gather_indices("gather_qmm", &x, w, lhs_indices, rhs_indices)?;
208    let mut inputs = vec![x.clone(), w.clone(), scales, biases, lhs, rhs.clone()];
209    if let Some(offsets) = sorted_offsets(&x, w, &rhs, lhs_indices.is_none() && sorted_indices)? {
210        inputs.push(offsets);
211    }
212    let op = crate::fast::FastOp::GatherQmm {
213        group_size,
214        bits,
215        transpose,
216    };
217    let kernel = transpose && crate::fast::on_gpu(dtype, x.size());
218    let out = if kernel {
219        Array::from_op(Primitive::Fast(op), inputs)?.swap_remove(0)
220    } else {
221        op.fallback(&inputs)?.swap_remove(0)
222    };
223    Ok(out)
224}
225
226/// The batch indices of a gathered product, broadcast together; either defaults to every
227/// batch of its operand, in order.
228fn gather_indices(
229    op: &'static str,
230    a: &Array,
231    b: &Array,
232    lhs: Option<&Array>,
233    rhs: Option<&Array>,
234) -> Result<(Array, Array)> {
235    let indices = |x: &Array, given: Option<&Array>| -> Result<Array> {
236        let batch = &x.shape()[..x.ndim() - 2];
237        match given {
238            Some(i)
239                if i.dtype().category() == Category::Float || i.dtype() == crate::DType::Bool =>
240            {
241                Err(Error::InvalidArgument(format!(
242                    "{op}: indices must be integers"
243                )))
244            }
245            Some(i) => astype(i, crate::DType::U32),
246            None => {
247                let n: usize = batch.iter().product();
248                let all = crate::arange(0.0, n as f64, 1.0, crate::DType::U32)?;
249                all.reshape(&batch.iter().map(|&d| d as isize).collect::<Vec<_>>())
250            }
251        }
252    };
253    let (l, r) = (indices(a, lhs)?, indices(b, rhs)?);
254    let batch = view::broadcast_shapes(op, l.shape(), r.shape())?;
255    Ok((broadcast_to(&l, &batch)?, broadcast_to(&r, &batch)?))
256}
257
258/// For rows sorted by weight batch (`a` is `[rows, 1, K]`, one row per output batch), where
259/// each batch's rows start: the rows per batch (scatter-added ones), then their exclusive
260/// running sum (MLX's grouped kernels walk every weight batch once).
261fn sorted_offsets(a: &Array, b: &Array, rhs: &Array, sorted: bool) -> Result<Option<Array>> {
262    let m = a.shape()[a.ndim() - 2];
263    if !(sorted && m == 1 && a.ndim() == 3 && rhs.ndim() == 1) {
264        return Ok(None);
265    }
266    let experts: usize = b.shape()[..b.ndim() - 2].iter().product();
267    let ones = crate::ones(rhs.shape(), crate::DType::I32)?;
268    let zeros = crate::zeros(&[experts], crate::DType::I32)?;
269    let counts = crate::scatter_add_along_axis(&zeros, rhs, &ones, 0)?;
270    Ok(Some(crate::cumsum(&counts, Some(0), false, false)?))
271}
272
273/// Matrix products of gathered batches, as MLX's `gather_mm` (mixture-of-experts layers in
274/// full precision): output batch `i` is `a[lhs_indices[i]] · b[rhs_indices[i]]` over the
275/// flattened batch dimensions of `a` (`[..., M, K]`) and `b` (`[..., K, N]`), with the
276/// indices broadcast together and either defaulting to every batch. `sorted_indices`
277/// promises sorted `rhs_indices`, which lets one kernel walk each batch of `b` once.
278///
279/// ```
280/// use nuvec_core::{Array, Result, gather_mm};
281///
282/// # fn main() -> Result<()> {
283/// // Two 2×2 "experts": the identity and twice the identity.
284/// let b = Array::from_slice(&[1.0f32, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0], &[2, 2, 2])?;
285/// let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 1, 2])?; // two rows
286/// let experts = Array::from_slice(&[1u32, 0], &[2])?;
287/// let y = gather_mm(&a, &b, None, Some(&experts), false)?;
288/// assert_eq!(y.to_vec::<f32>()?, [2.0, 4.0, 3.0, 4.0]);
289/// # Ok(())
290/// # }
291/// ```
292pub fn gather_mm(
293    a: &Array,
294    b: &Array,
295    lhs_indices: Option<&Array>,
296    rhs_indices: Option<&Array>,
297    sorted_indices: bool,
298) -> Result<Array> {
299    check_inputs(&[a, b])?;
300    let dtype = promote_types(a.dtype(), b.dtype());
301    if dtype.category() != Category::Float {
302        return Err(Error::UnsupportedDType {
303            op: "gather_mm",
304            dtype,
305        });
306    }
307    if a.ndim() < 2 || b.ndim() < 2 || a.shape()[a.ndim() - 1] != b.shape()[b.ndim() - 2] {
308        return Err(Error::InvalidArgument(format!(
309            "gather_mm: a must be [..., M, K] and b [..., K, N], got {:?} and {:?}",
310            a.shape(),
311            b.shape()
312        )));
313    }
314    let (a, b) = (astype(a, dtype)?, astype(b, dtype)?);
315    let (lhs, rhs) = gather_indices("gather_mm", &a, &b, lhs_indices, rhs_indices)?;
316    let mut inputs = vec![a.clone(), b.clone(), lhs, rhs.clone()];
317    if let Some(offsets) = sorted_offsets(&a, &b, &rhs, lhs_indices.is_none() && sorted_indices)? {
318        inputs.push(offsets);
319    }
320    Ok(crate::fast::record_or_fallback(crate::fast::FastOp::GatherMm, inputs)?.swap_remove(0))
321}
nuvec_core/quantized.rs

nuvec_core/
quantized.rs
