cherenkov/
quant.rs

1//! MLX affine quantization, CPU side. Layout (bits=4, group_size=64):
2//! - weight: u32 [out, in/8], 8 nibbles per word, LSB-first along the input dim
3//! - scales, biases: bf16 [out, in/64], one (scale, bias) per 64-element group
4//! - dequant: x = scale * q + bias with q in 0..=15
5
6use half::bf16;
7use rayon::prelude::*;
8
9pub const GROUP_SIZE: usize = 64;
10const NIBBLES_PER_WORD: usize = 8;
11
12/// Zero-copy view of one quantized linear layer's tensors.
13pub struct QLinear<'a> {
14    pub out_dim: usize,
15    pub in_dim: usize,
16    pub weight: &'a [u32],
17    pub scales: &'a [bf16],
18    pub biases: &'a [bf16],
19}
20
21impl<'a> QLinear<'a> {
22    /// Dequantize one output row into `dst` (len in_dim). Used for embedding
23    /// lookup and for testing.
24    pub fn dequant_row(&self, row: usize, dst: &mut [f32]) {
25        assert_eq!(dst.len(), self.in_dim);
26
27        let words_per_row = self.in_dim / NIBBLES_PER_WORD;
28        let groups_per_row = self.in_dim / GROUP_SIZE;
29        let words = &self.weight[row * words_per_row..(row + 1) * words_per_row];
30        let scales = &self.scales[row * groups_per_row..(row + 1) * groups_per_row];
31        let biases = &self.biases[row * groups_per_row..(row + 1) * groups_per_row];
32
33        for (wi, &word) in words.iter().enumerate() {
34            let base = wi * NIBBLES_PER_WORD;
35            let g = base / GROUP_SIZE;
36            let scale = scales[g].to_f32();
37            let bias = biases[g].to_f32();
38
39            for j in 0..NIBBLES_PER_WORD {
40                let q = (word >> (4 * j)) & 0xF;
41                dst[base + j] = scale * q as f32 + bias;
42            }
43        }
44    }
45
46    /// y = W x (f32 accumulate), parallel over output rows.
47    /// Grouped form: `y[r] = sum_g scale[r,g] * dot(q[r,g], x_g) + bias[r,g] * sum(x_g)`.
48    pub fn matvec(&self, x: &[f32], y: &mut [f32]) {
49        assert_eq!(x.len(), self.in_dim);
50        assert_eq!(y.len(), self.out_dim);
51
52        let groups_per_row = self.in_dim / GROUP_SIZE;
53        // Per-group plain sums of x, shared across all rows.
54        let xsums: Vec<f32> = x
55            .as_chunks::<GROUP_SIZE>()
56            .0
57            .iter()
58            .map(|g| g.iter().sum())
59            .collect();
60        let words_per_row = self.in_dim / NIBBLES_PER_WORD;
61
62        y.par_iter_mut().enumerate().for_each(|(r, out)| {
63            let words = &self.weight[r * words_per_row..(r + 1) * words_per_row];
64            let scales = &self.scales[r * groups_per_row..(r + 1) * groups_per_row];
65            let biases = &self.biases[r * groups_per_row..(r + 1) * groups_per_row];
66            let mut acc = 0.0f32;
67
68            const WORDS_PER_GROUP: usize = GROUP_SIZE / NIBBLES_PER_WORD;
69
70            for g in 0..groups_per_row {
71                let mut qdot = 0.0f32;
72                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
73
74                for wi in 0..WORDS_PER_GROUP {
75                    let word = words[g * WORDS_PER_GROUP + wi];
76                    let xw = &xg[wi * NIBBLES_PER_WORD..(wi + 1) * NIBBLES_PER_WORD];
77
78                    for (j, &value) in xw.iter().enumerate() {
79                        let q = (word >> (4 * j)) & 0xF;
80                        qdot += q as f32 * value;
81                    }
82                }
83
84                acc += scales[g].to_f32() * qdot + biases[g].to_f32() * xsums[g];
85            }
86
87            *out = acc;
88        });
89    }
90}
91
92#[cfg(test)]
93#[path = "../tests/unit/quant.rs"]
94mod tests;
cherenkov/quant.rs

cherenkov/
quant.rs
