nabled_embeddings/
quantize.rs

1//! int8 row quantization for embedding matrices (encode/decode + dequantized scoring).
2//!
3//! This is an honest, minimal quantization layer: it shrinks an `f32` embedding matrix to one
4//! `i8` per element plus one `f32` scale per row, and scores quantized matrices by **dequantizing
5//! back to `f32` and reusing the existing kernels**. There is no native int8 kernel and no new
6//! `nabled-linalg` surface; the win is ~4x smaller storage/transfer, traded against a small,
7//! bounded precision loss (each value is rounded to one of 255 levels scaled per row).
8//!
9//! # Scheme
10//!
11//! Per-row **symmetric** quantization. For row `r` with max absolute value `amax`:
12//! `scale[r] = amax / 127`, and each element is `round(x / scale[r])` clamped to `[-127, 127]`.
13//! A fully-zero row gets `scale = 0` and all-zero codes, and dequantizes back to zeros. The signed
14//! range is kept symmetric (`-127..=127`, leaving `-128` unused) so that `dequantize(quantize(x))`
15//! has no sign bias.
16
17use ndarray::{Array1, Array2, ArrayView2};
18
19use crate::error::EmbeddingError;
20use crate::similarity::{Metric, query_corpus_scores};
21
22/// The largest magnitude code used by the symmetric int8 scheme.
23const MAX_CODE: f32 = 127.0;
24
25/// A row-quantized `f32` matrix: one `i8` code per element plus one `f32` scale per row.
26#[derive(Debug, Clone, PartialEq)]
27pub struct QuantizedMatrix {
28    data:   Array2<i8>,
29    scales: Array1<f32>,
30}
31
32impl QuantizedMatrix {
33    /// The `(n_rows, n_cols)` int8 code matrix.
34    #[must_use]
35    pub fn data(&self) -> &Array2<i8> { &self.data }
36
37    /// The per-row `f32` scales (`len == n_rows`).
38    #[must_use]
39    pub fn scales(&self) -> &Array1<f32> { &self.scales }
40
41    /// Number of rows.
42    #[must_use]
43    pub fn nrows(&self) -> usize { self.data.nrows() }
44
45    /// Number of columns.
46    #[must_use]
47    pub fn ncols(&self) -> usize { self.data.ncols() }
48
49    /// Reconstruct (decode) the approximate `f32` matrix as `code * scale[row]`.
50    #[must_use]
51    pub fn dequantize(&self) -> Array2<f32> { dequantize(self) }
52
53    /// Construct a [`QuantizedMatrix`] from existing codes and scales (for example, decoded from a
54    /// wire format).
55    ///
56    /// # Errors
57    /// Returns [`EmbeddingError::EmptyInput`] when `data` is empty and
58    /// [`EmbeddingError::DimensionMismatch`] when `scales.len() != data.nrows()`.
59    pub fn from_parts(data: Array2<i8>, scales: Array1<f32>) -> Result<Self, EmbeddingError> {
60        if data.is_empty() {
61            return Err(EmbeddingError::EmptyInput);
62        }
63        if scales.len() != data.nrows() {
64            return Err(EmbeddingError::DimensionMismatch);
65        }
66        Ok(Self { data, scales })
67    }
68
69    /// Score this quantized matrix as queries against a quantized `corpus` under `metric`.
70    ///
71    /// Implemented as dequantize-then-existing-kernel: both operands are decoded to `f32` and
72    /// scored with [`query_corpus_scores`]. Results approximate the full-precision path within the
73    /// quantization tolerance.
74    ///
75    /// # Errors
76    /// Returns the same errors as [`query_corpus_scores`] (empty input, dimension mismatch, or
77    /// cosine zero-norm rows).
78    pub fn query_corpus_scores_quantized(
79        &self,
80        corpus: &QuantizedMatrix,
81        metric: Metric,
82    ) -> Result<Array2<f32>, EmbeddingError> {
83        let queries = self.dequantize();
84        let corpus = corpus.dequantize();
85        query_corpus_scores(&queries, &corpus, metric)
86    }
87}
88
89/// Quantize each row of `rows` to int8 with a per-row symmetric scale.
90///
91/// # Errors
92/// Returns [`EmbeddingError::EmptyInput`] when `rows` is empty.
93// Codes are clamped to `[-127, 127]` before the `as i8` cast, so truncation/sign-loss cannot occur
94// for the values produced here.
95#[allow(clippy::cast_possible_truncation)]
96pub fn quantize_rows(rows: &ArrayView2<'_, f32>) -> Result<QuantizedMatrix, EmbeddingError> {
97    if rows.is_empty() {
98        return Err(EmbeddingError::EmptyInput);
99    }
100    let (n_rows, n_cols) = rows.dim();
101    let mut data = Array2::<i8>::zeros((n_rows, n_cols));
102    let mut scales = Array1::<f32>::zeros(n_rows);
103
104    for (r, row) in rows.outer_iter().enumerate() {
105        let amax = row.iter().fold(0.0_f32, |acc, &v| acc.max(v.abs()));
106        if amax == 0.0 {
107            // Zero row: scale stays 0, codes stay 0.
108            continue;
109        }
110        let scale = amax / MAX_CODE;
111        scales[r] = scale;
112        for (c, &value) in row.iter().enumerate() {
113            let code = (value / scale).round().clamp(-MAX_CODE, MAX_CODE);
114            data[[r, c]] = code as i8;
115        }
116    }
117
118    Ok(QuantizedMatrix { data, scales })
119}
120
121/// Decode a [`QuantizedMatrix`] back to its approximate `f32` values (`code * scale[row]`).
122#[must_use]
123pub fn dequantize(matrix: &QuantizedMatrix) -> Array2<f32> {
124    let mut out = Array2::<f32>::zeros(matrix.data.dim());
125    out.outer_iter_mut().zip(matrix.data.outer_iter()).zip(matrix.scales.iter()).for_each(
126        |((mut out_row, code_row), &scale)| {
127            out_row.iter_mut().zip(code_row.iter()).for_each(|(value, &code)| {
128                *value = f32::from(code) * scale;
129            });
130        },
131    );
132    out
133}
134
135#[cfg(test)]
136mod tests {
137    use ndarray::arr2;
138
139    use super::*;
140
141    #[test]
142    fn round_trip_is_within_scale_tolerance() {
143        let rows = arr2(&[[1.0_f32, -2.0, 0.5], [0.1, 0.2, -0.3]]);
144        let q = quantize_rows(&rows.view()).unwrap();
145        let restored = q.dequantize();
146        for (r, row) in rows.outer_iter().enumerate() {
147            // Per-row reconstruction error is bounded by half a quantization step (scale/2).
148            let half_step = q.scales()[r] / 2.0 + 1e-6;
149            for (c, &orig) in row.iter().enumerate() {
150                assert!(
151                    (restored[[r, c]] - orig).abs() <= half_step,
152                    "row {r} col {c}: {orig} vs {}",
153                    restored[[r, c]]
154                );
155            }
156        }
157    }
158
159    #[test]
160    fn extreme_value_maps_to_max_code() {
161        let rows = arr2(&[[3.0_f32, -3.0, 1.5]]);
162        let q = quantize_rows(&rows.view()).unwrap();
163        // The two extremes map to +/-127.
164        assert_eq!(q.data()[[0, 0]], 127);
165        assert_eq!(q.data()[[0, 1]], -127);
166        assert!((q.scales()[0] - 3.0 / 127.0).abs() < 1e-9);
167    }
168
169    #[test]
170    fn zero_row_is_handled() {
171        let rows = arr2(&[[0.0_f32, 0.0], [1.0, 0.0]]);
172        let q = quantize_rows(&rows.view()).unwrap();
173        assert!(q.scales()[0].abs() < f32::EPSILON);
174        assert_eq!(q.data()[[0, 0]], 0);
175        assert_eq!(q.data()[[0, 1]], 0);
176        let restored = q.dequantize();
177        assert!(restored[[0, 0]].abs() < f32::EPSILON);
178        assert!(restored[[0, 1]].abs() < f32::EPSILON);
179    }
180
181    #[test]
182    fn dimensions_are_exposed() {
183        let rows = arr2(&[[1.0_f32, 2.0, 3.0], [4.0, 5.0, 6.0]]);
184        let q = quantize_rows(&rows.view()).unwrap();
185        assert_eq!(q.nrows(), 2);
186        assert_eq!(q.ncols(), 3);
187        assert_eq!(q.scales().len(), 2);
188    }
189
190    #[test]
191    fn quantize_rejects_empty() {
192        let rows = Array2::<f32>::zeros((0, 4));
193        assert_eq!(quantize_rows(&rows.view()).err(), Some(EmbeddingError::EmptyInput));
194    }
195
196    #[test]
197    fn from_parts_validates_shape() {
198        let data = Array2::<i8>::zeros((2, 3));
199        let scales = Array1::<f32>::zeros(3);
200        assert_eq!(
201            QuantizedMatrix::from_parts(data, scales).err(),
202            Some(EmbeddingError::DimensionMismatch)
203        );
204    }
205
206    #[test]
207    fn from_parts_rejects_empty() {
208        let data = Array2::<i8>::zeros((0, 0));
209        let scales = Array1::<f32>::zeros(0);
210        assert_eq!(
211            QuantizedMatrix::from_parts(data, scales).err(),
212            Some(EmbeddingError::EmptyInput)
213        );
214    }
215
216    #[test]
217    fn quantized_scoring_is_close_to_f32_path() {
218        let queries = arr2(&[[0.8_f32, 0.1, 0.5], [0.2, 0.9, 0.4]]);
219        let corpus = arr2(&[[0.7_f32, 0.2, 0.3], [0.1, 0.8, 0.6], [0.5, 0.5, 0.5]]);
220        let qq = quantize_rows(&queries.view()).unwrap();
221        let qc = quantize_rows(&corpus.view()).unwrap();
222        for metric in [Metric::Cosine, Metric::Dot, Metric::L2] {
223            let exact = query_corpus_scores(&queries, &corpus, metric).unwrap();
224            let approx = qq.query_corpus_scores_quantized(&qc, metric).unwrap();
225            for (lhs, rhs) in exact.iter().zip(approx.iter()) {
226                assert!((lhs - rhs).abs() < 5e-2, "{metric:?}: exact {lhs} vs quantized {rhs}");
227            }
228        }
229    }
230}
nabled_embeddings/quantize.rs

nabled_embeddings/
quantize.rs
