1use crate::error::CodecError;
18
19#[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 TernaryPacked = 3,
33 TernarySimd = 4,
35 TurboQuant4b = 5,
36 Sq8 = 6,
37 Pq = 7,
38 Itq3S = 8,
40 PolarQuant = 9,
42}
43
44impl QuantMode {
45 fn bits_per_weight(self) -> u32 {
48 match self {
49 QuantMode::Binary => 1,
50 QuantMode::RaBitQ => 1,
51 QuantMode::TernaryPacked => 2, 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#[repr(C)]
69#[derive(Debug, Clone, Copy)]
70pub struct QuantHeader {
71 pub quant_mode: u16,
73 pub dim: u16,
74 pub global_scale: f32,
76 pub residual_norm: f32,
78 pub dot_quantized: f32,
80 pub outlier_bitmask: u64,
83 pub reserved: [u8; 8],
85}
86
87const QUANT_HEADER_BYTES: usize = 32;
89
90const _: () = assert!(core::mem::size_of::<QuantHeader>() == QUANT_HEADER_BYTES);
91
92impl QuantHeader {
93 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 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
133const OUTLIER_ENTRY_BYTES: usize = 8;
137
138const ALIGN: usize = 128;
140
141pub 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#[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
169pub struct UnifiedQuantizedVector {
181 header: QuantHeader,
183 buf: Vec<u8>,
185 packed_bits_len: usize,
187}
188
189impl UnifiedQuantizedVector {
190 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 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 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 #[inline]
259 pub fn header(&self) -> &QuantHeader {
260 &self.header
261 }
262
263 #[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 #[inline]
272 pub fn outlier_count(&self) -> u32 {
273 self.header().outlier_bitmask.count_ones()
274 }
275
276 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 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 #[inline]
306 pub fn as_bytes(&self) -> &[u8] {
307 &self.buf
308 }
309}
310
311pub struct UnifiedQuantizedVectorRef<'a> {
317 header: QuantHeader,
318 buf: &'a [u8],
319 packed_bits_len: usize,
320}
321
322impl<'a> UnifiedQuantizedVectorRef<'a> {
323 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 #[inline]
359 pub fn header(&self) -> &QuantHeader {
360 &self.header
361 }
362
363 #[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 #[inline]
372 pub fn outlier_count(&self) -> u32 {
373 self.header().outlier_bitmask.count_ones()
374 }
375
376 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#[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]; 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 let bitmask: u64 = 1 << 5;
490 let header = make_header(QuantMode::Sq8, 64, bitmask);
491 let packed = vec![0u8; 64]; 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 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]; 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]; 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 let bitmask = header.outlier_bitmask;
582 let count = bitmask.count_ones() as usize;
583 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 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 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 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}