Skip to main content

nodedb_codec/vector_quant/
layout.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Unified Quantized Vector layout — the cache-aligned superset format that
4//! absorbs binary / ternary (BitNet 1.58) / 4-bit scalar / residual codecs
5//! without polymorphic indirection in the hot path.
6//!
7//! 128-byte alignment matches AVX-512 cache-line pair and avoids false
8//! sharing under thread-per-core execution.
9//!
10//! ## Outlier bitmask limit
11//!
12//! The `outlier_bitmask` in [`QuantHeader`] supports up to 64 outlier
13//! dimensions per vector. Callers that need to mark more than 64 outliers
14//! must bucket by 64-dim windows; multi-window support is out of scope for
15//! this module.
16
17use crate::error::CodecError;
18
19// ── QuantMode ──────────────────────────────────────────────────────────────
20
21/// Quantization mode discriminator stored in the header.
22///
23/// Values are **stable on disk** — never reorder, only append.
24#[repr(u16)]
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26#[non_exhaustive]
27pub enum QuantMode {
28    Binary = 0,
29    RaBitQ = 1,
30    Bbq = 2,
31    /// 5 trits/byte cold layout, 1.6 bpw.
32    TernaryPacked = 3,
33    /// 2 bpw expanded for SIMD hot path.
34    TernarySimd = 4,
35    TurboQuant4b = 5,
36    Sq8 = 6,
37    Pq = 7,
38    /// Reserved discriminant for ITQ3_S (Interleaved Ternary Quantization).
39    Itq3S = 8,
40    /// Reserved discriminant for PolarQuant (Cartesian→polar 0.5 bpw).
41    PolarQuant = 9,
42}
43
44impl QuantMode {
45    /// Bits per weight for each mode — used by [`target_size`] to compute the
46    /// packed-bits byte count.
47    fn bits_per_weight(self) -> u32 {
48        match self {
49            QuantMode::Binary => 1,
50            QuantMode::RaBitQ => 1,
51            QuantMode::TernaryPacked => 2, // 5 trits/byte ≈ 1.6 bpw; ceil to 2 for alignment
52            QuantMode::TernarySimd => 2,
53            QuantMode::Bbq => 1,
54            QuantMode::TurboQuant4b => 4,
55            QuantMode::Sq8 => 8,
56            QuantMode::Pq => 8,
57            QuantMode::Itq3S => 2,
58            QuantMode::PolarQuant => 4,
59        }
60    }
61}
62
63// ── QuantHeader ────────────────────────────────────────────────────────────
64
65/// 32-byte little-endian header preceding the packed bit array.
66///
67/// The wire layout is explicit and does not depend on Rust struct layout.
68#[repr(C)]
69#[derive(Debug, Clone, Copy)]
70pub struct QuantHeader {
71    /// [`QuantMode`] discriminant.
72    pub quant_mode: u16,
73    pub dim: u16,
74    /// BitNet absmean / SQ8 scale / RaBitQ rotation magnitude.
75    pub global_scale: f32,
76    /// RaBitQ ‖v−c‖ / TurboQuant ‖r‖ / BBQ centroid distance.
77    pub residual_norm: f32,
78    /// RaBitQ ⟨v, q(v)⟩ / BBQ corrective / TurboQuant QJL bias.
79    pub dot_quantized: f32,
80    /// Sparse outlier index — branchless gather via popcnt.
81    /// Supports up to 64 outlier dimensions per vector.
82    pub outlier_bitmask: u64,
83    /// QJL projection seed / OSAQ Hessian coeff / future use.
84    pub reserved: [u8; 8],
85}
86
87/// Serialized version-1 header width.
88const QUANT_HEADER_BYTES: usize = 32;
89
90const _: () = assert!(core::mem::size_of::<QuantHeader>() == QUANT_HEADER_BYTES);
91
92impl QuantHeader {
93    /// Serialize this header in the fixed version-1 little-endian layout.
94    fn to_le_bytes(self) -> [u8; QUANT_HEADER_BYTES] {
95        let mut bytes = [0u8; QUANT_HEADER_BYTES];
96        bytes[0..2].copy_from_slice(&self.quant_mode.to_le_bytes());
97        bytes[2..4].copy_from_slice(&self.dim.to_le_bytes());
98        bytes[4..8].copy_from_slice(&self.global_scale.to_le_bytes());
99        bytes[8..12].copy_from_slice(&self.residual_norm.to_le_bytes());
100        bytes[12..16].copy_from_slice(&self.dot_quantized.to_le_bytes());
101        bytes[16..24].copy_from_slice(&self.outlier_bitmask.to_le_bytes());
102        bytes[24..32].copy_from_slice(&self.reserved);
103        bytes
104    }
105
106    /// Parse a fixed version-1 little-endian header without alignment assumptions.
107    fn from_le_bytes(bytes: &[u8]) -> Result<Self, CodecError> {
108        let header = bytes.get(..QUANT_HEADER_BYTES).ok_or_else(|| CodecError::LayoutError {
109            detail: format!(
110                "buffer too short for quantized vector header: need {QUANT_HEADER_BYTES} bytes, got {}",
111                bytes.len()
112            ),
113        })?;
114        Ok(Self {
115            quant_mode: u16::from_le_bytes([header[0], header[1]]),
116            dim: u16::from_le_bytes([header[2], header[3]]),
117            global_scale: f32::from_le_bytes([header[4], header[5], header[6], header[7]]),
118            residual_norm: f32::from_le_bytes([header[8], header[9], header[10], header[11]]),
119            dot_quantized: f32::from_le_bytes([header[12], header[13], header[14], header[15]]),
120            outlier_bitmask: u64::from_le_bytes([
121                header[16], header[17], header[18], header[19], header[20], header[21], header[22],
122                header[23],
123            ]),
124            reserved: header[24..32]
125                .try_into()
126                .map_err(|_| CodecError::LayoutError {
127                    detail: "quantized vector reserved header bytes are truncated".into(),
128                })?,
129        })
130    }
131}
132
133// ── Layout constants ────────────────────────────────────────────────────────
134
135/// Byte size of one outlier entry: (dim_index: u32, value: f32).
136const OUTLIER_ENTRY_BYTES: usize = 8;
137
138/// Cache-line pair alignment — also the minimum allocation unit.
139const ALIGN: usize = 128;
140
141// ── target_size ─────────────────────────────────────────────────────────────
142
143/// Compute the buffer size (rounded up to a 128-byte multiple) required to
144/// hold a [`UnifiedQuantizedVector`] with the given parameters.
145///
146/// This is the pre-sizing helper for callers that want to allocate before
147/// constructing.
148pub fn target_size(quant_mode: QuantMode, dim: u16, outlier_count: u32) -> usize {
149    let packed_bits_bytes = packed_bits_len(quant_mode, dim);
150    let outlier_bytes = outlier_count as usize * OUTLIER_ENTRY_BYTES;
151    let raw = QUANT_HEADER_BYTES + packed_bits_bytes + outlier_bytes;
152    round_up_128(raw)
153}
154
155// ── Internal helpers ────────────────────────────────────────────────────────
156
157#[inline]
158fn packed_bits_len(quant_mode: QuantMode, dim: u16) -> usize {
159    let bpw = quant_mode.bits_per_weight() as usize;
160    let total_bits = dim as usize * bpw;
161    total_bits.div_ceil(8)
162}
163
164#[inline]
165fn round_up_128(n: usize) -> usize {
166    (n + ALIGN - 1) & !(ALIGN - 1)
167}
168
169// ── UnifiedQuantizedVector ──────────────────────────────────────────────────
170
171/// Owned, 128-byte-aligned unified quantized vector buffer.
172///
173/// Layout (contiguous bytes):
174/// ```text
175/// [ QuantHeader (32 B) | packed_bits (variable) | outlier_payload (8 B × n) | tail_pad ]
176/// ```
177///
178/// The total allocation is always a multiple of 128 bytes (one AVX-512
179/// cache-line pair).
180pub struct UnifiedQuantizedVector {
181    /// Parsed fixed-layout header, retained separately from the byte buffer.
182    header: QuantHeader,
183    /// Backing storage. Always a multiple of 128 bytes.
184    buf: Vec<u8>,
185    /// Byte length of the packed-bits region (excludes header and outliers).
186    packed_bits_len: usize,
187}
188
189impl UnifiedQuantizedVector {
190    /// Construct from an explicit header, packed-bit slice, and sparse
191    /// outlier list.
192    ///
193    /// `outliers` is a slice of `(dim_index, value)` pairs.  The
194    /// `outlier_bitmask` in `header` must have exactly one bit set for each
195    /// entry in `outliers`, and bits must correspond to entries in ascending
196    /// `dim_index` order (i.e. popcnt-dense order).
197    ///
198    /// # Errors
199    ///
200    /// Returns [`CodecError::LayoutError`] if:
201    /// - `outliers.len()` does not match `popcnt(header.outlier_bitmask)`.
202    /// - Any `dim_index` in `outliers` is ≥ 64 (bitmask only covers 64 dims).
203    pub fn new(
204        header: QuantHeader,
205        packed_bits: &[u8],
206        outliers: &[(u32, f32)],
207    ) -> Result<Self, CodecError> {
208        let expected_outlier_count = header.outlier_bitmask.count_ones() as usize;
209        if outliers.len() != expected_outlier_count {
210            return Err(CodecError::LayoutError {
211                detail: format!(
212                    "outlier count mismatch: bitmask has {} bits set but {} outliers provided",
213                    expected_outlier_count,
214                    outliers.len()
215                ),
216            });
217        }
218        for &(dim_idx, _) in outliers {
219            if dim_idx >= 64 {
220                return Err(CodecError::LayoutError {
221                    detail: format!("outlier dim_index {dim_idx} exceeds bitmask capacity of 64"),
222                });
223            }
224        }
225
226        let header_bytes = QUANT_HEADER_BYTES;
227        let outlier_bytes = outliers.len() * OUTLIER_ENTRY_BYTES;
228        let raw = header_bytes + packed_bits.len() + outlier_bytes;
229        let total = round_up_128(raw);
230
231        let mut buf = vec![0u8; total];
232
233        buf[..header_bytes].copy_from_slice(&header.to_le_bytes());
234
235        // Write packed bits.
236        let pb_start = header_bytes;
237        let pb_end = pb_start + packed_bits.len();
238        buf[pb_start..pb_end].copy_from_slice(packed_bits);
239
240        // Write outlier payload: each entry is u32 dim_index || f32 value (LE).
241        let mut off = pb_end;
242        for &(dim_idx, value) in outliers {
243            buf[off..off + 4].copy_from_slice(&dim_idx.to_le_bytes());
244            buf[off + 4..off + 8].copy_from_slice(&value.to_le_bytes());
245            off += OUTLIER_ENTRY_BYTES;
246        }
247
248        Ok(Self {
249            header,
250            buf,
251            packed_bits_len: packed_bits.len(),
252        })
253    }
254
255    // ── Accessors ────────────────────────────────────────────────────────────
256
257    /// Parsed fixed-layout header.
258    #[inline]
259    pub fn header(&self) -> &QuantHeader {
260        &self.header
261    }
262
263    /// Slice of the packed-bit region.
264    #[inline]
265    pub fn packed_bits(&self) -> &[u8] {
266        let start = QUANT_HEADER_BYTES;
267        &self.buf[start..start + self.packed_bits_len]
268    }
269
270    /// Number of outlier entries, computed via popcnt of `outlier_bitmask`.
271    #[inline]
272    pub fn outlier_count(&self) -> u32 {
273        self.header().outlier_bitmask.count_ones()
274    }
275
276    /// Return the outlier `(dim_index, value)` for the dimension at position
277    /// `slot` in the bitmask.
278    ///
279    /// `slot` is the dimension index (0–63).  Returns `None` if the bit for
280    /// `slot` is not set in `outlier_bitmask`, or if `slot ≥ 64`.
281    ///
282    /// Uses a branchless popcnt to find the dense offset into the outlier
283    /// payload.
284    pub fn outlier_at(&self, slot: u32) -> Option<(u32, f32)> {
285        if slot >= 64 {
286            return None;
287        }
288        let bitmask = self.header().outlier_bitmask;
289        if bitmask & (1u64 << slot) == 0 {
290            return None;
291        }
292        // Number of set bits below `slot` gives the dense array index.
293        let mask = bitmask & ((1u64 << slot).wrapping_sub(1));
294        let offset = mask.count_ones() as usize;
295
296        let header_bytes = QUANT_HEADER_BYTES;
297        let base = header_bytes + self.packed_bits_len + offset * OUTLIER_ENTRY_BYTES;
298
299        let dim_idx = u32::from_le_bytes(self.buf[base..base + 4].try_into().ok()?);
300        let value = f32::from_le_bytes(self.buf[base + 4..base + 8].try_into().ok()?);
301        Some((dim_idx, value))
302    }
303
304    /// Full backing buffer suitable for direct I/O.
305    #[inline]
306    pub fn as_bytes(&self) -> &[u8] {
307        &self.buf
308    }
309}
310
311// ── UnifiedQuantizedVectorRef ───────────────────────────────────────────────
312
313/// Zero-copy borrowed view into a `UnifiedQuantizedVector` buffer.
314///
315/// Suitable for reads from the io_uring page cache — no allocation required.
316pub struct UnifiedQuantizedVectorRef<'a> {
317    header: QuantHeader,
318    buf: &'a [u8],
319    packed_bits_len: usize,
320}
321
322impl<'a> UnifiedQuantizedVectorRef<'a> {
323    /// Borrow a raw byte slice as a quantized vector view.
324    ///
325    /// # Errors
326    ///
327    /// Returns [`CodecError::LayoutError`] if the slice is shorter than 32 bytes
328    /// (minimum header size).
329    pub fn from_bytes(buf: &'a [u8], packed_bits_len: usize) -> Result<Self, CodecError> {
330        let header = QuantHeader::from_le_bytes(buf)?;
331        let outlier_bytes = (header.outlier_bitmask.count_ones() as usize)
332            .checked_mul(OUTLIER_ENTRY_BYTES)
333            .ok_or_else(|| CodecError::LayoutError {
334                detail: "quantized vector outlier length overflow".into(),
335            })?;
336        let required = QUANT_HEADER_BYTES
337            .checked_add(packed_bits_len)
338            .and_then(|length| length.checked_add(outlier_bytes))
339            .ok_or_else(|| CodecError::LayoutError {
340                detail: "quantized vector length overflow".into(),
341            })?;
342        if buf.len() < required {
343            return Err(CodecError::LayoutError {
344                detail: format!(
345                    "buffer too short: need at least {required} bytes, got {}",
346                    buf.len()
347                ),
348            });
349        }
350        Ok(Self {
351            header,
352            buf,
353            packed_bits_len,
354        })
355    }
356
357    /// Parsed fixed-layout header.
358    #[inline]
359    pub fn header(&self) -> &QuantHeader {
360        &self.header
361    }
362
363    /// Packed-bits slice.
364    #[inline]
365    pub fn packed_bits(&self) -> &[u8] {
366        let start = QUANT_HEADER_BYTES;
367        &self.buf[start..start + self.packed_bits_len]
368    }
369
370    /// Number of outlier entries via popcnt.
371    #[inline]
372    pub fn outlier_count(&self) -> u32 {
373        self.header().outlier_bitmask.count_ones()
374    }
375
376    /// Outlier lookup — same semantics as [`UnifiedQuantizedVector::outlier_at`].
377    pub fn outlier_at(&self, slot: u32) -> Option<(u32, f32)> {
378        if slot >= 64 {
379            return None;
380        }
381        let bitmask = self.header().outlier_bitmask;
382        if bitmask & (1u64 << slot) == 0 {
383            return None;
384        }
385        let mask = bitmask & ((1u64 << slot).wrapping_sub(1));
386        let offset = mask.count_ones() as usize;
387
388        let base = QUANT_HEADER_BYTES
389            .checked_add(self.packed_bits_len)?
390            .checked_add(offset.checked_mul(OUTLIER_ENTRY_BYTES)?)?;
391        let entry = self.buf.get(base..base.checked_add(OUTLIER_ENTRY_BYTES)?)?;
392
393        let dim_idx = u32::from_le_bytes(entry[0..4].try_into().ok()?);
394        let value = f32::from_le_bytes(entry[4..8].try_into().ok()?);
395        Some((dim_idx, value))
396    }
397}
398
399// ── Tests ───────────────────────────────────────────────────────────────────
400
401#[cfg(test)]
402mod tests {
403    use super::*;
404
405    fn make_header(mode: QuantMode, dim: u16, bitmask: u64) -> QuantHeader {
406        QuantHeader {
407            quant_mode: mode as u16,
408            dim,
409            global_scale: 1.5,
410            residual_norm: 0.25,
411            dot_quantized: 2.5,
412            outlier_bitmask: bitmask,
413            reserved: [0xAB; 8],
414        }
415    }
416
417    #[test]
418    fn header_is_32_bytes() {
419        assert_eq!(QUANT_HEADER_BYTES, 32);
420    }
421
422    #[test]
423    fn header_parses_explicit_little_endian_fixture_from_unaligned_subslice() {
424        let header = QuantHeader {
425            quant_mode: 0x1234,
426            dim: 0x5678,
427            global_scale: f32::from_bits(0x3FC0_0000),
428            residual_norm: f32::from_bits(0xBF40_0000),
429            dot_quantized: f32::from_bits(0x4040_0000),
430            outlier_bitmask: 0x1122_3344_5566_7788,
431            reserved: [1, 2, 3, 4, 5, 6, 7, 8],
432        };
433        let mut fixture = vec![0xFF];
434        fixture.extend_from_slice(&header.to_le_bytes());
435        fixture.extend_from_slice(&[0u8; 64 * OUTLIER_ENTRY_BYTES]);
436
437        let view = UnifiedQuantizedVectorRef::from_bytes(&fixture[1..], 0)
438            .expect("unaligned header fixture parses");
439        let parsed = view.header();
440        assert_eq!(parsed.quant_mode, 0x1234);
441        assert_eq!(parsed.dim, 0x5678);
442        assert_eq!(parsed.global_scale.to_bits(), 0x3FC0_0000);
443        assert_eq!(parsed.residual_norm.to_bits(), 0xBF40_0000);
444        assert_eq!(parsed.dot_quantized.to_bits(), 0x4040_0000);
445        assert_eq!(parsed.outlier_bitmask, 0x1122_3344_5566_7788);
446        assert_eq!(parsed.reserved, [1, 2, 3, 4, 5, 6, 7, 8]);
447    }
448
449    #[test]
450    fn target_size_is_128_multiple() {
451        for mode in [
452            QuantMode::Binary,
453            QuantMode::RaBitQ,
454            QuantMode::TernarySimd,
455            QuantMode::TurboQuant4b,
456            QuantMode::Sq8,
457        ] {
458            for dim in [64u16, 128, 256, 512, 1536] {
459                for outliers in [0u32, 1, 8, 64] {
460                    let sz = target_size(mode, dim, outliers);
461                    assert_eq!(
462                        sz % 128,
463                        0,
464                        "target_size not 128-aligned for {mode:?}/{dim}/{outliers}"
465                    );
466                    assert!(
467                        sz >= 128,
468                        "target_size below minimum for {mode:?}/{dim}/{outliers}"
469                    );
470                }
471            }
472        }
473    }
474
475    #[test]
476    fn no_outliers_roundtrip() {
477        let header = make_header(QuantMode::Binary, 128, 0);
478        let packed = vec![0xFFu8; 16]; // 128 dims / 8 = 16 bytes
479        let vec = UnifiedQuantizedVector::new(header, &packed, &[]).unwrap();
480
481        assert_eq!(vec.outlier_count(), 0);
482        assert_eq!(vec.packed_bits(), packed.as_slice());
483        assert_eq!(vec.as_bytes().len() % 128, 0);
484    }
485
486    #[test]
487    fn one_outlier_roundtrip() {
488        // Bit 5 set → dim_index 5 is an outlier.
489        let bitmask: u64 = 1 << 5;
490        let header = make_header(QuantMode::Sq8, 64, bitmask);
491        let packed = vec![0u8; 64]; // 64 dims × 8bpw = 64 bytes
492        let outliers = [(5u32, 42.0f32)];
493        let vec = UnifiedQuantizedVector::new(header, &packed, &outliers).unwrap();
494
495        assert_eq!(vec.outlier_count(), 1);
496        let (dim, val) = vec.outlier_at(5).expect("bit 5 should be set");
497        assert_eq!(dim, 5);
498        assert!((val - 42.0).abs() < f32::EPSILON);
499        assert!(vec.outlier_at(0).is_none());
500        assert!(vec.outlier_at(6).is_none());
501    }
502
503    #[test]
504    fn eight_outliers_roundtrip() {
505        // Bits 0,3,7,12,20,33,50,63 set.
506        let bits: &[u32] = &[0, 3, 7, 12, 20, 33, 50, 63];
507        let mut bitmask: u64 = 0;
508        for &b in bits {
509            bitmask |= 1 << b;
510        }
511        let header = make_header(QuantMode::TurboQuant4b, 128, bitmask);
512        let packed = vec![0xAAu8; 64]; // 128 dims × 4bpw = 64 bytes
513        let outlier_list: Vec<(u32, f32)> = bits
514            .iter()
515            .enumerate()
516            .map(|(i, &b)| (b, i as f32 * 1.1))
517            .collect();
518        let vec = UnifiedQuantizedVector::new(header, &packed, &outlier_list).unwrap();
519
520        assert_eq!(vec.outlier_count(), 8);
521        for (i, &b) in bits.iter().enumerate() {
522            let (dim, val) = vec
523                .outlier_at(b)
524                .unwrap_or_else(|| panic!("outlier at {b} missing"));
525            assert_eq!(dim, b);
526            assert!(
527                (val - i as f32 * 1.1f32).abs() < 1e-5,
528                "value mismatch at dim {b}"
529            );
530        }
531    }
532
533    #[test]
534    fn as_bytes_reborrow_via_ref() {
535        let bitmask: u64 = 1 << 10;
536        let header = make_header(QuantMode::RaBitQ, 64, bitmask);
537        let packed = vec![0u8; 8]; // 64 dims / 8 = 8 bytes (1 bpw)
538        let outliers = [(10u32, 7.77f32)];
539        let vec = UnifiedQuantizedVector::new(header, &packed, &outliers).unwrap();
540
541        let bytes = vec.as_bytes();
542        let packed_bits_len = vec.packed_bits_len;
543        let vref = UnifiedQuantizedVectorRef::from_bytes(bytes, packed_bits_len).unwrap();
544
545        assert_eq!(vref.outlier_count(), 1);
546        let (dim, val) = vref.outlier_at(10).unwrap();
547        assert_eq!(dim, 10);
548        assert!((val - 7.77).abs() < 1e-5);
549    }
550
551    #[test]
552    fn borrowed_view_rejects_truncated_outliers_and_length_overflow() {
553        let header = QuantHeader {
554            quant_mode: QuantMode::Bbq as u16,
555            dim: 1,
556            global_scale: 1.0,
557            residual_norm: 0.0,
558            dot_quantized: 0.0,
559            outlier_bitmask: 1,
560            reserved: [0; 8],
561        };
562        let bytes = header.to_le_bytes();
563        assert!(UnifiedQuantizedVectorRef::from_bytes(&bytes, 0).is_err());
564        assert!(UnifiedQuantizedVectorRef::from_bytes(&bytes, usize::MAX).is_err());
565    }
566
567    #[test]
568    fn header_field_roundtrip() {
569        let header = QuantHeader {
570            quant_mode: QuantMode::Bbq as u16,
571            dim: 512,
572            global_scale: 4.5,
573            residual_norm: 0.99,
574            dot_quantized: -1.23,
575            outlier_bitmask: 0xDEAD_BEEF_0000_0001,
576            reserved: [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08],
577        };
578        let packed = vec![0u8; packed_bits_len(QuantMode::Bbq, 512)];
579
580        // Build bitmask consistent with header.
581        let bitmask = header.outlier_bitmask;
582        let count = bitmask.count_ones() as usize;
583        // Generate dummy outlier entries for each set bit (lowest bits first).
584        let mut outliers: Vec<(u32, f32)> = Vec::with_capacity(count);
585        for bit in 0u32..64 {
586            if bitmask & (1u64 << bit) != 0 {
587                outliers.push((bit, bit as f32));
588            }
589        }
590
591        let vec = UnifiedQuantizedVector::new(header, &packed, &outliers).unwrap();
592        let h = vec.header();
593
594        assert_eq!(h.quant_mode, QuantMode::Bbq as u16);
595        assert_eq!(h.dim, 512);
596        assert!((h.global_scale - 4.5).abs() < 1e-5);
597        assert!((h.residual_norm - 0.99).abs() < 1e-5);
598        assert!((h.dot_quantized - (-1.23)).abs() < 1e-5);
599        assert_eq!(h.outlier_bitmask, 0xDEAD_BEEF_0000_0001);
600        assert_eq!(h.reserved, [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
601    }
602
603    #[test]
604    fn outlier_ordering_popcnt() {
605        // Outliers at dims [3, 17, 40] — bit 17 should be the second entry.
606        let bitmask: u64 = (1 << 3) | (1 << 17) | (1 << 40);
607        let header = make_header(QuantMode::Sq8, 64, bitmask);
608        let packed = vec![0u8; 64];
609        let outliers = [(3u32, 100.0f32), (17u32, 200.0f32), (40u32, 300.0f32)];
610        let vec = UnifiedQuantizedVector::new(header, &packed, &outliers).unwrap();
611
612        // outlier_at(17) should return the second entry (value 200.0).
613        let (dim, val) = vec.outlier_at(17).expect("dim 17 should be an outlier");
614        assert_eq!(dim, 17);
615        assert!((val - 200.0).abs() < f32::EPSILON);
616
617        let (dim0, val0) = vec.outlier_at(3).expect("dim 3 should be an outlier");
618        assert_eq!(dim0, 3);
619        assert!((val0 - 100.0).abs() < f32::EPSILON);
620
621        let (dim2, val2) = vec.outlier_at(40).expect("dim 40 should be an outlier");
622        assert_eq!(dim2, 40);
623        assert!((val2 - 300.0).abs() < f32::EPSILON);
624    }
625
626    #[test]
627    fn out_of_range_slot_returns_none() {
628        let header = make_header(QuantMode::Binary, 64, 0);
629        let packed = vec![0u8; 8];
630        let vec = UnifiedQuantizedVector::new(header, &packed, &[]).unwrap();
631
632        assert!(vec.outlier_at(64).is_none(), "slot 64 is out of range");
633        assert!(vec.outlier_at(80).is_none(), "slot 80 is out of range");
634        assert!(
635            vec.outlier_at(u32::MAX).is_none(),
636            "slot u32::MAX is out of range"
637        );
638    }
639
640    #[test]
641    fn outlier_count_mismatch_is_error() {
642        // Bitmask says 1 outlier but we provide 0.
643        let bitmask: u64 = 1 << 2;
644        let header = make_header(QuantMode::Binary, 64, bitmask);
645        let packed = vec![0u8; 8];
646        let err = UnifiedQuantizedVector::new(header, &packed, &[]);
647        assert!(
648            err.is_err(),
649            "should fail when outlier count mismatches bitmask"
650        );
651    }
652}