zeppelin_embed/quant/
bits4.rs

1//! Four-bit Extended-RaBitQ row encoding.
2//!
3//! The conceptual codebook is the normalized half-integer grid
4//! `{-7.5, -6.5, ..., -0.5, 0.5, ..., 6.5, 7.5}^d`. It is fixed by the bit
5//! width: there is no learned codebook, training set, or first-query setup.
6//! Random rotation is optional in Extended-RaBitQ and is off by default in v1;
7//! the recall harness must validate that identity-rotation assumption for each
8//! embedding model.
9//!
10//! For input `v`, Algorithm 1 selects `y` maximizing `<y,v>/||y||`. At scale
11//! zero each signed magnitude is `0.5`. Coordinate `i` then crosses its seven
12//! rounding boundaries at `t = level / |v_i|`, `level=1..=7`. Sorting all
13//! critical values, applying equal thresholds as one group, and retaining the
14//! best normalized projection enumerates every realizable rounded direction
15//! exactly in `O(2^B d log d)` time.
16//!
17//! The stored RaBitQ correction is `c = ||v||^2 / <y,v>`. The paper writes the
18//! estimator with normalized `y`; because `||y||` occurs in numerator and
19//! denominator it cancels, leaving `c * <y,q_hat>`. The companion norm is kept
20//! for reconstruction/error bounds and future distance bookkeeping.
21
22use crate::kernels::{MAX_DOT_I8_DIMENSION, dot_bit4_prepared, score_bit4_prepared_batch};
23
24use super::QuantError;
25
26mod controlled;
27pub(crate) use controlled::{
28    Bit4ControlError, Bit4Scratch, est_dot_bit4_controlled, prepare_bit4_query_controlled,
29    quantize_bit4_controlled,
30};
31
32const CODES_PER_BYTE: usize = 2;
33const MAX_MAGNITUDE_LEVEL: u8 = 7;
34
35/// Per-row scalars required by the four-bit unbiased dot estimator.
36#[derive(Clone, Copy, Debug, PartialEq)]
37#[repr(C)]
38pub struct Bit4Factors {
39    scale: f32,
40    normalized_norm: f32,
41    normalized_correction: f32,
42}
43
44impl Bit4Factors {
45    /// Reconstructs the permanent three-f32 persisted factor record.
46    #[must_use]
47    pub const fn from_persisted(
48        scale: f32,
49        normalized_norm: f32,
50        normalized_correction: f32,
51    ) -> Self {
52        Self {
53            scale,
54            normalized_norm,
55            normalized_correction,
56        }
57    }
58
59    /// Returns the permanent little-endian record fields in declaration order.
60    #[must_use]
61    pub const fn persisted_fields(self) -> [f32; 3] {
62        [self.scale, self.normalized_norm, self.normalized_correction]
63    }
64
65    /// Returns the original row's Euclidean norm.
66    ///
67    /// A finite `f32` row scale and finite normalized `f32` norm are expanded
68    /// in `f64`, so the logical norm remains usable above `f32::MAX` without
69    /// enlarging the three-scalar factor record.
70    #[must_use]
71    pub const fn norm(self) -> f64 {
72        self.scale as f64 * self.normalized_norm as f64
73    }
74
75    /// Returns `||v||^2 / <y, v>` for the selected grid direction.
76    ///
77    /// The stored normalized correction is multiplied by the finite row scale
78    /// in `f64`, avoiding overflow in persisted metadata.
79    #[must_use]
80    pub const fn correction(self) -> f64 {
81        self.scale as f64 * self.normalized_correction as f64
82    }
83
84    /// Returns the L2 error bound for norm-scaled code reconstruction.
85    ///
86    /// Every selected grid coordinate is sign-aligned with the input, hence
87    /// the direction cosine is non-negative and the deterministic bound is
88    /// `sqrt(2) ||v||`. The exact grid search normally gives a tighter error;
89    /// retaining this conservative bound avoids a fourth per-row scalar.
90    #[must_use]
91    pub fn reconstruction_error_bound(self) -> f64 {
92        std::f64::consts::SQRT_2 * self.norm()
93    }
94
95    pub(crate) const fn scoring_parts(self) -> (f32, f32) {
96        (self.scale, self.normalized_correction)
97    }
98}
99
100/// Query-side signed-byte representation prepared once and reused per row.
101///
102/// Coordinates are symmetrically quantized to `[-127, 127]` with seeded
103/// stochastic rounding, so their reconstruction is unbiased in expectation.
104/// The private kernel representation groups even then odd coordinates within
105/// each 32-value block, matching the packed row's high/low nibble vectors.
106#[derive(Clone, Debug, PartialEq)]
107pub struct Bit4Query {
108    /// Even coordinates followed by odd coordinates in each 32-value block.
109    codes: Vec<i8>,
110    code_sum: i32,
111    scale_half: f64,
112}
113
114impl Bit4Query {
115    /// Returns the query coordinate count.
116    #[must_use]
117    pub fn len(&self) -> usize {
118        self.codes.len()
119    }
120
121    /// Returns whether the query has zero coordinates.
122    #[must_use]
123    pub fn is_empty(&self) -> bool {
124        self.codes.is_empty()
125    }
126
127    pub(crate) fn resident_bytes(&self) -> usize {
128        self.codes.capacity()
129    }
130
131    /// Returns the prepared code bytes, pre-interleave sum, and half-scale
132    /// exactly as consumed by the scoring kernel.
133    #[cfg(any(test, feature = "test-seams"))]
134    #[doc(hidden)]
135    #[must_use]
136    pub fn observation_parts(&self) -> (&[i8], i32, f64) {
137        (&self.codes, self.code_sum, self.scale_half)
138    }
139
140    pub(crate) fn kernel_parts(&self) -> (&[i8], i32, f64) {
141        (&self.codes, self.code_sum, self.scale_half)
142    }
143}
144
145#[derive(Clone, Copy, Debug)]
146struct SplitMix64 {
147    state: u64,
148}
149
150impl SplitMix64 {
151    const fn new(seed: u64) -> Self {
152        Self { state: seed }
153    }
154
155    fn next_open_unit_f64(&mut self) -> f64 {
156        self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15);
157        let mut value = self.state;
158        value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
159        value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
160        value ^= value >> 31;
161        ((value >> 11) as f64 + 0.5) * (1.0 / 9_007_199_254_740_992.0)
162    }
163}
164
165#[derive(Clone, Copy, Debug)]
166struct CriticalValue {
167    threshold: f64,
168    coordinate: usize,
169    level: u8,
170    magnitude: f64,
171}
172
173/// Quantizes one row to four-bit Extended-RaBitQ codes.
174///
175/// Two coordinates occupy each byte in coordinate order, high nibble first.
176/// Unsigned code `u` represents `u - 7.5`. For an odd dimension, the unused
177/// low nibble in the final byte is canonical zero padding.
178///
179/// # Errors
180///
181/// Returns [`QuantError::EmptyVector`] for an empty input,
182/// [`QuantError::DimensionTooLarge`] above the i8 scoring-kernel limit,
183/// [`QuantError::NonFinite`] for NaN or infinity, and
184/// [`QuantError::OutputLength`] unless `out.len() == ceil(v.len() / 2)`.
185/// Validation completes before `out` is modified.
186pub fn quantize_bit4(v: &[f32], out: &mut [u8]) -> Result<Bit4Factors, QuantError> {
187    validate_input(v, out.len())?;
188    let mut scratch = Bit4Scratch::compatibility(v.len());
189    match quantize_bit4_controlled(v, out, &mut scratch, &mut |_| Ok::<_, ()>(())) {
190        Ok(factors) => Ok(factors),
191        Err(Bit4ControlError::Quant(error)) => Err(error),
192        Err(Bit4ControlError::Memory) => Err(QuantError::Allocation),
193        Err(Bit4ControlError::Control(())) => Err(QuantError::Allocation),
194    }
195}
196
197/// Reconstructs the norm-scaled selected four-bit code direction.
198///
199/// The result is `||v|| * y / ||y||`. Serving uses [`est_dot_bit4`] directly;
200/// reconstruction exists to specify and verify the codec's geometric error.
201///
202/// # Errors
203///
204/// Returns the code-length and padding errors documented by
205/// [`est_dot_bit4`]. The output length defines the logical dimension of the
206/// final partial byte.
207pub fn dequantize_bit4(
208    codes: &[u8],
209    factors: Bit4Factors,
210    out: &mut [f32],
211) -> Result<(), QuantError> {
212    let unpacked = unpack(codes, out.len())?;
213    if factors.scale == 0.0 {
214        out.fill(0.0);
215        return Ok(());
216    }
217    let doubled_code_norm = unpacked
218        .iter()
219        .map(|&value| {
220            let value = f64::from(value);
221            value * value
222        })
223        .sum::<f64>()
224        .sqrt();
225    for (output, &value) in out.iter_mut().zip(&unpacked) {
226        *output = finite_f32(factors.norm() * f64::from(value) / doubled_code_norm);
227    }
228    Ok(())
229}
230
231/// Prepares one query for repeated four-bit row scoring.
232///
233/// The seed controls reproducible stochastic rounding only. Random rotation is
234/// not performed by this v1 API. A zero query receives scale zero and all-zero
235/// signed codes.
236///
237/// # Errors
238///
239/// Returns [`QuantError::EmptyVector`], [`QuantError::DimensionTooLarge`], or
240/// [`QuantError::NonFinite`] under the same input policy as
241/// [`quantize_bit4`].
242pub fn prepare_bit4_query(q: &[f32], seed: u64) -> Result<Bit4Query, QuantError> {
243    #[cfg(any(test, feature = "test-seams"))]
244    super::QUERY_PREPARATIONS.with(|calls| {
245        if let Some(calls) = calls.borrow_mut().as_mut() {
246            calls.bit4.push((q.len(), seed));
247        }
248    });
249    match prepare_bit4_query_controlled(q, seed, &mut |_| Ok::<_, ()>(())) {
250        Ok((query, _)) => Ok(query),
251        Err(Bit4ControlError::Quant(error)) => Err(error),
252        Err(Bit4ControlError::Memory) => Err(QuantError::Allocation),
253        Err(Bit4ControlError::Control(())) => Err(QuantError::Allocation),
254    }
255}
256
257/// Estimates a dot product directly from one packed four-bit row.
258///
259/// The native packed kernel scores unsigned nibbles directly, then applies
260/// `dot(z, 2*u - 15) = 2*dot(z, u) - 15*sum(z)` once per row. The query's
261/// block-interleaved representation pairs directly with high- and low-nibble
262/// vectors, so scoring needs no expanded row, ZIP interleave, or per-vector
263/// grid conversion. If the prepared query reconstructs as `scale*z`, the
264/// integer kernel gives `<y,q_hat> = scale * dot(2*y,z) / 2`; the stored
265/// correction then yields the Extended-RaBitQ estimate.
266///
267/// For an odd dimension, the unused low nibble must be canonical zero padding.
268///
269/// # Errors
270///
271/// Returns [`QuantError::CodeLength`] for a dimension mismatch and
272/// [`QuantError::NonZeroPadding`] for non-canonical trailing fields.
273pub fn est_dot_bit4(
274    query: &Bit4Query,
275    codes: &[u8],
276    factors: Bit4Factors,
277) -> Result<f32, QuantError> {
278    match est_dot_bit4_controlled(query, codes, factors, &mut |_| {
279        Ok::<_, std::convert::Infallible>(())
280    }) {
281        Ok(value) => Ok(value),
282        Err(Bit4ControlError::Quant(error)) => Err(error),
283        Err(Bit4ControlError::Memory) => Err(QuantError::Allocation),
284        Err(Bit4ControlError::Control(never)) => match never {},
285    }
286}
287
288/// Estimates dot products for contiguous packed four-bit rows in one dispatch.
289///
290/// `factors` and `out` contain one entry per row. Row validation is performed
291/// before entering the dispatched microkernel, and the scalar table remains
292/// the bit-exact behavioral oracle on machines without the selected SIMD tier.
293///
294/// # Errors
295///
296/// Returns [`QuantError::OutputLength`] for a factor/output count mismatch,
297/// [`QuantError::CodeLength`] for a batch byte-length mismatch, or
298/// [`QuantError::NonZeroPadding`] for a non-canonical odd-dimension row.
299pub fn est_dot_bit4_batch(
300    query: &Bit4Query,
301    codes: &[u8],
302    factors: &[Bit4Factors],
303    out: &mut [f32],
304) -> Result<(), QuantError> {
305    if out.len() != factors.len() {
306        return Err(QuantError::OutputLength {
307            expected: factors.len(),
308            actual: out.len(),
309        });
310    }
311    let row_bytes = query.codes.len().div_ceil(CODES_PER_BYTE);
312    let expected = row_bytes.saturating_mul(factors.len());
313    if codes.len() != expected {
314        return Err(QuantError::CodeLength {
315            expected,
316            actual: codes.len(),
317        });
318    }
319    if !query.codes.len().is_multiple_of(CODES_PER_BYTE) {
320        for row in codes.chunks_exact(row_bytes) {
321            validate_padding(row, query.codes.len())?;
322        }
323    }
324    score_bit4_prepared_batch(
325        &query.codes,
326        query.code_sum,
327        query.scale_half,
328        codes,
329        query.codes.len(),
330        factors,
331        out,
332    );
333    Ok(())
334}
335
336fn validate_input(v: &[f32], output_len: usize) -> Result<(), QuantError> {
337    validate_vector(v)?;
338    let expected = v.len().div_ceil(CODES_PER_BYTE);
339    if output_len != expected {
340        return Err(QuantError::OutputLength {
341            expected,
342            actual: output_len,
343        });
344    }
345    Ok(())
346}
347
348fn validate_vector_shape(v: &[f32]) -> Result<(), QuantError> {
349    if v.is_empty() {
350        return Err(QuantError::EmptyVector);
351    }
352    if v.len() > MAX_DOT_I8_DIMENSION {
353        return Err(QuantError::DimensionTooLarge {
354            actual: v.len(),
355            maximum: MAX_DOT_I8_DIMENSION,
356        });
357    }
358    Ok(())
359}
360
361fn validate_vector(v: &[f32]) -> Result<(), QuantError> {
362    validate_vector_shape(v)?;
363    if let Some((index, _)) = v.iter().enumerate().find(|(_, value)| !value.is_finite()) {
364        return Err(QuantError::NonFinite { index });
365    }
366    Ok(())
367}
368
369fn unpack(codes: &[u8], dimension: usize) -> Result<Vec<i8>, QuantError> {
370    validate_code(codes, dimension)?;
371
372    let mut unpacked = Vec::with_capacity(dimension);
373    for &byte in codes {
374        for shift in [4_u32, 0] {
375            if unpacked.len() == dimension {
376                break;
377            }
378            let unsigned = (byte >> shift) & 0b1111;
379            unpacked.push((i16::from(unsigned) * 2 - 15) as i8);
380        }
381    }
382    Ok(unpacked)
383}
384
385fn validate_code(codes: &[u8], dimension: usize) -> Result<(), QuantError> {
386    let expected = dimension.div_ceil(CODES_PER_BYTE);
387    if codes.len() != expected {
388        return Err(QuantError::CodeLength {
389            expected,
390            actual: codes.len(),
391        });
392    }
393    validate_padding(codes, dimension)?;
394    Ok(())
395}
396
397fn validate_padding(codes: &[u8], dimension: usize) -> Result<(), QuantError> {
398    if dimension.is_multiple_of(CODES_PER_BYTE) {
399        return Ok(());
400    }
401    let mask = 0x0f;
402    if let Some(&byte) = codes.last()
403        && byte & mask != 0
404    {
405        return Err(QuantError::NonZeroPadding { byte, mask });
406    }
407    Ok(())
408}
409
410fn normalized_score_squared(numerator: f64, norm_squared: f64) -> f64 {
411    numerator * numerator / norm_squared
412}
413
414fn finite_f32(value: f64) -> f32 {
415    value.clamp(f64::from(f32::MIN), f64::from(f32::MAX)) as f32
416}
417
418fn apply_event(event: &CriticalValue, numerator: &mut f64, norm_squared: &mut f64) {
419    *numerator += event.magnitude;
420    *norm_squared += 2.0 * f64::from(event.level);
421}
zeppelin_embed/quant/bits4.rs

zeppelin_embed/quant/
bits4.rs
