From 57756e69ec5f0c44d4cddfca6aec3011879365b8 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 19 Sep 2026 20:35:24 -0700 Subject: [PATCH] feat(ann): optional int8 storage for the index's vector copy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The HNSW index keeps its own copy of every vector, which at 100K x 384 f32 is ~146 MiB — the largest single item in the 2.43x footprint now that the agent stores embeddings once. `Storage::Int8` cuts that copy to a quarter by scaling each row to i8. The scale is per row, not global. Unit-length rows in d dimensions have components around 1/sqrt(d), so a fixed [-1, 1] scale spends fewer than 12 of the 255 levels on a 128-dimensional vector; measured against an exact ranking that gives 0.35 top-10 overlap. Scaling each row by its own largest component uses the full range and brings it to 0.99. Quantised distances still cost recall on their own, and `ef` does not buy it back because the loss is in the distances rather than the graph: at N=100K recall@10 tops out at 0.967 against f32's 0.9995. Re-scoring a wider candidate pool against the exact vectors removes the gap (0.9940 vs 0.9945 at ef=64) for ~13% of query throughput and ~16% of build time. That is the intended use, so it is what the test asserts — against ground truth, not against the f32 index, whose own mistakes a re-scored search is entitled to get right. Default is unchanged: `Storage::Float32`, chosen by every existing constructor. Serialized indexes carry f32 vectors and no storage tag, so a quantised index is rebuilt rather than loaded; `compact()` keeps the storage it was given. The harness grows `--int8` and `--rerank` axes, and reports the storage in each table header. Co-Authored-By: Claude Opus 5 (1M context) --- crates/clawhdf5-ann/src/hnsw.rs | 493 +++++++++++++++--- crates/clawhdf5-ann/src/lib.rs | 2 +- .../clawhdf5-bench/src/bin/search_harness.rs | 56 +- 3 files changed, 483 insertions(+), 68 deletions(-) diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index eb7124d..a4104e4 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -154,6 +154,237 @@ impl Ord for FarCandidate { } } +/// How the index keeps its copy of the vectors. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum Storage { + /// Exactly as given: `dim * 4` bytes per vector. + #[default] + Float32, + /// Each component scaled to an `i8`: `dim` bytes per vector, a quarter of + /// the space, at some cost in precision. + /// + /// Only meaningful for [`DistanceMetric::Cosine`]: rows are stored + /// unit-length, so a quantised dot product reconstructs the similarity + /// directly. Requesting it for `L2` keeps `Float32`, because an L2 + /// distance cannot be recovered from a dot product alone. + Int8, +} + +/// The index's copy of the vectors, flat and row-major. +#[derive(Debug, Clone)] +enum Vectors { + F32 { + dim: usize, + flat: Vec, + }, + /// `flat[i * dim + j]` is component `j` of vector `i` divided by + /// `scales[i]`; multiplying back recovers it. + /// + /// The scale is per row rather than global. A unit-length row in `d` + /// dimensions has components around `1/sqrt(d)`, so a fixed `[-1, 1]` + /// scale spends fewer than 12 of the 255 levels on a 128-dimensional + /// vector and the reconstruction error swamps the gaps between near + /// neighbours — measured at 0.35 top-10 overlap with the exact ranking. + /// Scaling each row by its own largest component uses the full range. + Int8 { + dim: usize, + flat: Vec, + scales: Vec, + }, +} + +/// Levels either side of zero. 127, not 128, so the range is symmetric. +const INT8_LEVELS: f32 = 127.0; + +/// Quantise one row, returning the codes and the scale that inverts them. +fn quantise_row(v: &[f32], out: &mut Vec) -> f32 { + let max_abs = v.iter().fold(0.0f32, |m, x| m.max(x.abs())); + if max_abs <= f32::MIN_POSITIVE { + out.extend(core::iter::repeat_n(0i8, v.len())); + return 0.0; + } + let inv = INT8_LEVELS / max_abs; + out.extend( + v.iter() + .map(|x| (x * inv).round().clamp(-INT8_LEVELS, INT8_LEVELS) as i8), + ); + max_abs / INT8_LEVELS +} + +impl Vectors { + fn new(dim: usize, storage: Storage, metric: DistanceMetric) -> Self { + match storage { + Storage::Int8 if metric == DistanceMetric::Cosine => Vectors::Int8 { + dim, + flat: Vec::new(), + scales: Vec::new(), + }, + _ => Vectors::F32 { + dim, + flat: Vec::new(), + }, + } + } + + fn dim(&self) -> usize { + match self { + Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim, + } + } + + fn storage(&self) -> Storage { + match self { + Vectors::F32 { .. } => Storage::Float32, + Vectors::Int8 { .. } => Storage::Int8, + } + } + + fn len(&self) -> usize { + let dim = self.dim(); + if dim == 0 { + return 0; + } + match self { + Vectors::F32 { flat, .. } => flat.len() / dim, + Vectors::Int8 { flat, .. } => flat.len() / dim, + } + } + + /// Set the row width, for a store seeded empty by `new`. + fn set_dim(&mut self, new_dim: usize) { + match self { + Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim = new_dim, + } + } + + fn push(&mut self, vector: &[f32]) { + match self { + Vectors::F32 { flat, .. } => flat.extend_from_slice(vector), + Vectors::Int8 { flat, scales, .. } => scales.push(quantise_row(vector, flat)), + } + } + + /// Row `i` as `f32`, for callers that need the values back (serialization, + /// and the f32 fast paths). Quantised rows are reconstructed, so this is + /// lossy in exactly the way the storage is. + fn row(&self, i: usize) -> Vec { + let dim = self.dim(); + let start = i * dim; + match self { + Vectors::F32 { flat, .. } => flat[start..start + dim].to_vec(), + Vectors::Int8 { flat, scales, .. } => flat[start..start + dim] + .iter() + .map(|&q| f32::from(q) * scales[i]) + .collect(), + } + } + + /// Distance between two stored vectors. + fn dist(&self, a: usize, b: usize, metric: DistanceMetric) -> f32 { + let dim = self.dim(); + match self { + Vectors::F32 { flat, .. } => { + let (x, y) = (a * dim, b * dim); + compute_distance(&flat[x..x + dim], &flat[y..y + dim], metric) + } + Vectors::Int8 { flat, scales, .. } => { + let (x, y) = (a * dim, b * dim); + let dot = dot_i8(&flat[x..x + dim], &flat[y..y + dim]); + 1.0 - dot as f32 * scales[a] * scales[b] + } + } + } + + /// Distance from a prepared query to stored vector `i`. + fn dist_query(&self, query: &Query, i: usize, metric: DistanceMetric) -> f32 { + let dim = self.dim(); + let start = i * dim; + match (self, query) { + (Vectors::F32 { flat, .. }, Query::F32(q)) => { + compute_distance(q, &flat[start..start + dim], metric) + } + (Vectors::Int8 { flat, scales, .. }, Query::Int8(q, q_scale)) => { + let dot = dot_i8(q, &flat[start..start + dim]); + 1.0 - dot as f32 * q_scale * scales[i] + } + // Mixed forms cannot occur: `Query` is built from the same storage. + _ => f32::MAX, + } + } + + /// Build a store from prepared rows. + fn from_rows(rows: &[Vec], storage: Storage, metric: DistanceMetric) -> Self { + let dim = rows.first().map_or(0, Vec::len); + let mut out = Vectors::new(dim, storage, metric); + for row in rows { + out.push(row); + } + out + } + + /// Prepare `query` for comparison against this store. + fn query(&self, query: Vec) -> Query { + match self { + Vectors::F32 { .. } => Query::F32(query), + Vectors::Int8 { .. } => { + let mut codes = Vec::with_capacity(query.len()); + let scale = quantise_row(&query, &mut codes); + Query::Int8(codes, scale) + } + } + } +} + +/// What a layer search is measuring distance *to*: an incoming query, or a +/// node already in the index (which is what insertion compares against). +enum Target<'a> { + Query(&'a Query), + Node(usize), +} + +impl Vectors { + fn dist_to(&self, target: &Target<'_>, i: usize, metric: DistanceMetric) -> f32 { + match target { + Target::Query(q) => self.dist_query(q, i, metric), + Target::Node(n) => self.dist(*n, i, metric), + } + } +} + +/// A search query in whichever form the store compares against. +enum Query { + F32(Vec), + /// Codes and the scale that inverts them, as in [`Vectors::Int8`]. + Int8(Vec, f32), +} + +/// Sum of products, widened so it cannot overflow: `dim` terms of at most +/// `127 * 127`, so `i32` suffices for any realistic dimension. +fn dot_i8(a: &[i8], b: &[i8]) -> i32 { + // Four independent accumulators over 32-lane blocks: the widening product + // has to sit in a fixed-length chunk for the vectoriser to see it, and the + // separate accumulators keep it off one dependency chain. + const LANE: usize = 8; + let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>(); + let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>(); + let mut acc = [0i32; 4]; + for (x, y) in a_blocks.iter().zip(b_blocks) { + for (lane, slot) in acc.iter_mut().enumerate() { + let mut sum = 0i32; + for k in 0..LANE { + sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]); + } + *slot += sum; + } + } + let tail: i32 = a_tail + .iter() + .zip(b_tail) + .map(|(&x, &y)| i32::from(x) * i32::from(y)) + .sum(); + acc[0] + acc[1] + acc[2] + acc[3] + tail +} + /// Magic for [`HnswIndex::graph_to_bytes`]. const GRAPH_MAGIC: &[u8; 4] = b"CHG1"; @@ -174,8 +405,8 @@ pub const HNSW_FORMAT_VERSION: i64 = 2; /// HDF5 format. #[derive(Debug, Clone)] pub struct HnswIndex { - /// All vectors in the index. - vectors: Vec>, + /// All vectors in the index, flat and row-major. + vectors: Vectors, /// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs. graph: Vec>>, /// Soft-deletion flags, one per node. Deleted nodes remain in the graph for @@ -214,6 +445,20 @@ impl HnswIndex { m: usize, ef_construction: usize, metric: DistanceMetric, + ) -> Self { + Self::build_with(vectors, m, ef_construction, metric, Storage::default()) + } + + /// Build an index, choosing how the vectors are stored. + /// + /// [`Storage::Int8`] keeps them at a quarter of the size; see its docs for + /// what that costs and when it applies. + pub fn build_with( + vectors: &[Vec], + m: usize, + ef_construction: usize, + metric: DistanceMetric, + storage: Storage, ) -> Self { assert!(!vectors.is_empty(), "cannot build index from empty vectors"); assert!(m >= 2, "m must be at least 2"); @@ -224,8 +469,11 @@ impl HnswIndex { let m_max0 = m * 2; let n = vectors.len(); - let prepared: Vec> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect(); - let vectors: &[Vec] = &prepared; + let mut prepared = Vectors::new(dim, storage, metric); + for v in vectors { + prepared.push(&prepare(v.clone(), metric)); + } + let vectors = &prepared; // Assign levels to all nodes let mut node_levels = Vec::with_capacity(n); @@ -322,9 +570,20 @@ impl HnswIndex { /// point for incremental [`HnswIndex::insert`] and as the result of /// [`HnswIndex::compact`] when every vector has been deleted. pub fn new(m: usize, ef_construction: usize, metric: DistanceMetric) -> Self { + Self::new_with(m, ef_construction, metric, Storage::default()) + } + + /// [`HnswIndex::new`], choosing how the vectors are stored. + pub fn new_with( + m: usize, + ef_construction: usize, + metric: DistanceMetric, + storage: Storage, + ) -> Self { assert!(m >= 2, "m must be at least 2"); Self { - vectors: Vec::new(), + // The dimension is set by the first insert. + vectors: Vectors::new(0, storage, metric), graph: Vec::new(), deleted: Vec::new(), entry_point: 0, @@ -351,7 +610,8 @@ impl HnswIndex { // Seed an empty index. if id == 0 { let node_level = assign_level(0, self.m); - self.vectors.push(vector); + self.vectors.set_dim(vector.len()); + self.vectors.push(&vector); self.deleted.push(false); self.node_levels.push(node_level); self.graph = (0..=node_level).map(|_| vec![Vec::new(); 1]).collect(); @@ -361,12 +621,12 @@ impl HnswIndex { assert_eq!( vector.len(), - self.vectors[0].len(), + self.vectors.dim(), "insert dimension mismatch" ); let node_level = assign_level(id, self.m); - self.vectors.push(vector); + self.vectors.push(&vector); self.deleted.push(false); self.node_levels.push(node_level); @@ -387,7 +647,7 @@ impl HnswIndex { ep = greedy_closest( &self.vectors, &self.graph[layer], - &self.vectors[id], + &Target::Node(id), ep, self.metric, ); @@ -400,7 +660,7 @@ impl HnswIndex { let neighbors = search_layer( &self.vectors, &self.graph[layer], - &self.vectors[id], + &Target::Node(id), ep, self.ef_construction, self.metric, @@ -466,16 +726,25 @@ impl HnswIndex { pub fn compact(&mut self) -> Vec> { let mut mapping = vec![None; self.vectors.len()]; let mut surviving: Vec> = Vec::with_capacity(self.active_len()); - for (old, v) in self.vectors.iter().enumerate() { + for (old, slot) in mapping.iter_mut().enumerate() { if !self.deleted[old] { - mapping[old] = Some(surviving.len()); - surviving.push(v.clone()); + *slot = Some(surviving.len()); + surviving.push(self.vectors.row(old)); } } + // Rebuilding must keep the storage the caller chose; a compaction is + // not the place to silently quadruple the index's memory. + let storage = self.vectors.storage(); *self = if surviving.is_empty() { - Self::new(self.m, self.ef_construction, self.metric) + Self::new_with(self.m, self.ef_construction, self.metric, storage) } else { - Self::build_with_metric(&surviving, self.m, self.ef_construction, self.metric) + Self::build_with( + &surviving, + self.m, + self.ef_construction, + self.metric, + storage, + ) }; mapping } @@ -490,24 +759,22 @@ impl HnswIndex { /// # Returns /// A vector of `(id, distance)` pairs sorted by distance (closest first). pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(usize, f32)> { - if self.vectors.is_empty() { + if self.vectors.len() == 0 { return Vec::new(); } - assert_eq!( - query.len(), - self.vectors[0].len(), - "query dimension mismatch" - ); + assert_eq!(query.len(), self.vectors.dim(), "query dimension mismatch"); let ef = ef.max(k); - let prepared_query = prepare(query.to_vec(), self.metric); - let query = prepared_query.as_slice(); + // Prepared and, for a quantised store, quantised once per search + // rather than once per comparison. + let prepared = self.vectors.query(prepare(query.to_vec(), self.metric)); + let target = Target::Query(&prepared); let mut ep = self.entry_point; let top_layer = self.graph.len().saturating_sub(1); // Greedy search from top layer down to layer 1 for layer in (1..=top_layer).rev() { - ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric); + ep = greedy_closest(&self.vectors, &self.graph[layer], &target, ep, self.metric); } // Search layer 0 for the ef nearest *live* nodes. Deleted nodes are @@ -516,7 +783,7 @@ impl HnswIndex { let candidates = search_layer( &self.vectors, &self.graph[0], - query, + &target, ep, ef, self.metric, @@ -543,14 +810,13 @@ impl HnswIndex { pub fn to_hdf5_bytes(&self) -> Result, FormatError> { let mut fw = FmtWriter::new(); let n = self.vectors.len(); - let dim = if n > 0 { self.vectors[0].len() } else { 0 }; + let dim = self.vectors.dim(); // Flatten vectors into a 1D array for storage - let flat_vectors: Vec = self - .vectors - .iter() - .flat_map(|v| v.iter().copied()) - .collect(); + let mut flat_vectors: Vec = Vec::with_capacity(n * dim); + for i in 0..n { + flat_vectors.extend_from_slice(&self.vectors.row(i)); + } let mut group = fw.create_group("ann"); @@ -709,7 +975,9 @@ impl HnswIndex { }; Ok(Self { - vectors, + // Serialized files carry f32 vectors and no storage tag: a + // quantised index is rebuilt, not loaded. + vectors: Vectors::from_rows(&vectors, Storage::Float32, metric), graph, deleted, entry_point, @@ -773,6 +1041,16 @@ impl HnswIndex { /// `bytes` is validated — a corrupt or mismatched graph is an error, never /// an index that panics or walks out of bounds during a search. pub fn from_graph_bytes(bytes: &[u8], vectors: Vec>) -> Result { + Self::from_graph_bytes_with(bytes, vectors, Storage::default()) + } + + /// As [`from_graph_bytes`](Self::from_graph_bytes), choosing how the + /// rehydrated vectors are stored. + pub fn from_graph_bytes_with( + bytes: &[u8], + vectors: Vec>, + storage: Storage, + ) -> Result { let bad = |what: &str| FormatError::SerializationError(format!("HNSW graph: {what}")); let body_len = bytes .len() @@ -863,7 +1141,14 @@ impl HnswIndex { } Ok(Self { - vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(), + vectors: Vectors::from_rows( + &vectors + .into_iter() + .map(|v| prepare(v, metric)) + .collect::>(), + storage, + metric, + ), graph, deleted, entry_point, @@ -882,16 +1167,17 @@ impl HnswIndex { /// Returns true if the index is empty. pub fn is_empty(&self) -> bool { - self.vectors.is_empty() + self.vectors.len() == 0 + } + + /// How this index stores its copy of the vectors. + pub fn storage(&self) -> Storage { + self.vectors.storage() } /// Returns the dimension of vectors in the index. pub fn dimension(&self) -> usize { - if self.vectors.is_empty() { - 0 - } else { - self.vectors[0].len() - } + self.vectors.dim() } /// Returns the number of layers in the graph. @@ -916,17 +1202,17 @@ impl HnswIndex { /// Greedy search: find the single closest node to `query` starting from `ep`. fn greedy_closest( - vectors: &[Vec], + vectors: &Vectors, layer: &[Vec], - query: &[f32], + target: &Target<'_>, mut ep: usize, metric: DistanceMetric, ) -> usize { - let mut best_dist = compute_distance(query, &vectors[ep], metric); + let mut best_dist = vectors.dist_to(target, ep, metric); loop { let mut changed = false; for &neighbor in &layer[ep] { - let d = compute_distance(query, &vectors[neighbor], metric); + let d = vectors.dist_to(target, neighbor, metric); if d < best_dist { best_dist = d; ep = neighbor; @@ -950,15 +1236,15 @@ fn greedy_closest( /// instead meant a query whose neighbourhood had been deleted got back fewer /// than `k` results, or none, however many live records were nearby. fn search_layer( - vectors: &[Vec], + vectors: &Vectors, layer: &[Vec], - query: &[f32], + target: &Target<'_>, ep: usize, ef: usize, metric: DistanceMetric, skip: Option<&[bool]>, ) -> Vec { - let ep_dist = compute_distance(query, &vectors[ep], metric); + let ep_dist = vectors.dist_to(target, ep, metric); // Min-heap of candidates to explore let mut candidates = BinaryHeap::new(); @@ -980,7 +1266,7 @@ fn search_layer( visited.begin(vectors.len()); visited.insert(ep); search_layer_visit( - vectors, layer, query, ef, metric, skip, visited, candidates, results, + vectors, layer, target, ef, metric, skip, visited, candidates, results, ) }) } @@ -1023,9 +1309,9 @@ thread_local! { #[allow(clippy::too_many_arguments)] fn search_layer_visit( - vectors: &[Vec], + vectors: &Vectors, layer: &[Vec], - query: &[f32], + target: &Target<'_>, ef: usize, metric: DistanceMetric, skip: Option<&[bool]>, @@ -1044,7 +1330,7 @@ fn search_layer_visit( continue; } - let d = compute_distance(query, &vectors[neighbor], metric); + let d = vectors.dist_to(target, neighbor, metric); let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance); if d < furthest_dist || results.len() < ef { @@ -1095,7 +1381,7 @@ fn search_layer_visit( /// remaining slots are then filled with the closest rejected candidates, so a /// node is never left under-connected. fn select_neighbors( - vectors: &[Vec], + vectors: &Vectors, candidates: &[(usize, f32)], max_conn: usize, metric: DistanceMetric, @@ -1111,7 +1397,7 @@ fn select_neighbors( } let diverse = selected .iter() - .all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node); + .all(|&s| vectors.dist(id, s, metric) > dist_to_node); if diverse { selected.push(id); } else { @@ -1136,7 +1422,7 @@ fn batch_len(linked: usize) -> usize { /// layers, found by searching the graph as it currently stands. #[allow(clippy::too_many_arguments)] fn plan_batch( - vectors: &[Vec], + vectors: &Vectors, graph: &[Vec>], node_levels: &[usize], batch: std::ops::Range, @@ -1150,7 +1436,7 @@ fn plan_batch( let mut ep = entry_point; // Phase 1: greedy descent from the top layer down to node_level + 1. for layer in (node_level + 1..=ep_level).rev() { - ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric); + ep = greedy_closest(vectors, &graph[layer], &Target::Node(i), ep, metric); } // Phase 2: search and select on every layer the node lives on. let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1); @@ -1159,7 +1445,7 @@ fn plan_batch( let neighbors = search_layer( vectors, &graph[layer], - &vectors[i], + &Target::Node(i), ep, ef_construction, metric, @@ -1186,7 +1472,7 @@ fn plan_batch( /// Prune every `(layer, node)` neighbour list in `overflowed` back to its /// limit. Each list belongs to a different node, so they are independent. fn prune_overflowed( - vectors: &[Vec], + vectors: &Vectors, graph: &mut [Vec>], overflowed: Vec<(usize, usize)>, (m, m_max0): (usize, usize), @@ -1226,7 +1512,7 @@ const PARALLEL_MIN: usize = 8; /// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16 /// cores; bulk builds batch their pruning instead, see `prune_overflowed`.) fn link_back( - vectors: &[Vec], + vectors: &Vectors, layer: &mut [Vec], new_id: usize, selected: &[usize], @@ -1248,7 +1534,7 @@ fn link_back( /// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`]. fn prune_connections( - vectors: &[Vec], + vectors: &Vectors, neighbors: &mut Vec, node: usize, max_conn: usize, @@ -1259,7 +1545,7 @@ fn prune_connections( } let mut scored: Vec<(usize, f32)> = neighbors .iter() - .map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric))) + .map(|&n| (n, vectors.dist(node, n, metric))) .collect(); scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0))); *neighbors = select_neighbors(vectors, &scored, max_conn, metric); @@ -1519,6 +1805,88 @@ mod tests { assert!(recall >= 0.95, "incremental recall@10 = {recall}"); } + #[test] + fn int8_storage_needs_an_exact_re_score_to_match_f32() { + // Cosine only: rows are unit-length, so a quantised dot product + // reconstructs the similarity directly. + // + // Not the `clustered` generator: its clusters are far tighter than any + // real embedding, so neighbours sit closer together than the + // quantisation error and top-10 identity there is noise — that would + // measure the fixture, not the storage. + let mut vectors = make_random_vectors(3060, 128, 5); + let queries = vectors.split_off(3000); + let f32_index = + HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Float32); + let quantised = + HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Int8); + assert_eq!(quantised.storage(), Storage::Int8); + + // Ground truth, not the f32 index's answers: re-scoring can beat that + // index, and measuring against it would score being right as drift. + let truth: Vec> = queries + .iter() + .map(|q| { + let mut d: Vec<(usize, f32)> = vectors + .iter() + .enumerate() + .map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::Cosine))) + .collect(); + d.sort_by(|a, b| a.1.total_cmp(&b.1)); + d[..10].iter().map(|x| x.0).collect() + }) + .collect(); + + let recall = |got: &dyn Fn(&[f32]) -> Vec| -> f64 { + let mut hits = 0; + for (q, want) in queries.iter().zip(&truth) { + hits += got(q).iter().filter(|id| want.contains(id)).count(); + } + hits as f64 / (10 * queries.len()) as f64 + }; + + let exact_recall = recall(&|q| f32_index.search(q, 10, 64).iter().map(|r| r.0).collect()); + let raw_recall = recall(&|q| quantised.search(q, 10, 64).iter().map(|r| r.0).collect()); + // Quantised distances alone cost recall, and `ef` cannot buy it back: + // the loss is in the distances, not in the graph. + assert!( + raw_recall < exact_recall, + "int8 alone should cost recall: {raw_recall} vs {exact_recall}" + ); + + // Re-scoring a wider candidate pool against the exact vectors — what a + // caller holding them (the agent's embedding cache) does — puts it + // back, because only the *ordering* was approximate. + let rescored_recall = recall(&|q| { + let mut pool: Vec<(usize, f32)> = quantised + .search(q, 40, 64) + .into_iter() + .map(|(id, _)| { + ( + id, + compute_distance(q, &vectors[id], DistanceMetric::Cosine), + ) + }) + .collect(); + pool.sort_by(|a, b| a.1.total_cmp(&b.1)); + pool.truncate(10); + pool.into_iter().map(|p| p.0).collect() + }); + assert!( + rescored_recall >= exact_recall - 0.01, + "int8 + exact re-score should match f32: {rescored_recall} vs {exact_recall} (raw {raw_recall})" + ); + } + + #[test] + fn int8_storage_falls_back_to_f32_for_non_cosine_metrics() { + // L2 distance is not recoverable from a quantised dot product, so the + // store silently stays f32 rather than returning wrong distances. + let vectors = clustered(100, 8, 5, 3); + let index = HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::L2, Storage::Int8); + assert_eq!(index.storage(), Storage::Float32); + } + #[test] fn deletions_near_the_query_do_not_shrink_or_degrade_results() { let mut vectors = clustered(2040, 16, 20, 11); @@ -1650,6 +2018,7 @@ mod tests { vec![1.2, 0.0], // 3 vec![-2.0, 0.0], // 4 ]; + let store = Vectors::from_rows(&vectors, Storage::Float32, DistanceMetric::L2); let scored: Vec<(usize, f32)> = (1..5) .map(|i| { ( @@ -1659,12 +2028,12 @@ mod tests { }) .collect(); assert_eq!( - select_neighbors(&vectors, &scored, 2, DistanceMetric::L2), + select_neighbors(&store, &scored, 2, DistanceMetric::L2), [1, 4] ); // Spare capacity is filled with the closest rejected candidates. assert_eq!( - select_neighbors(&vectors, &scored, 3, DistanceMetric::L2), + select_neighbors(&store, &scored, 3, DistanceMetric::L2), [1, 4, 2] ); } @@ -1790,7 +2159,7 @@ mod tests { // Verify vectors match for i in 0..loaded.len() { - assert_eq!(loaded.vectors[i], index.vectors[i]); + assert_eq!(loaded.vectors.row(i), index.vectors.row(i)); } } diff --git a/crates/clawhdf5-ann/src/lib.rs b/crates/clawhdf5-ann/src/lib.rs index bb5d042..f10470f 100644 --- a/crates/clawhdf5-ann/src/lib.rs +++ b/crates/clawhdf5-ann/src/lib.rs @@ -5,4 +5,4 @@ mod hnsw; -pub use hnsw::{DistanceMetric, HnswIndex}; +pub use hnsw::{DistanceMetric, HnswIndex, Storage}; diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index b8902a8..411d148 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -24,7 +24,7 @@ use std::time::{Duration, Instant}; use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry}; -use clawhdf5_ann::{DistanceMetric, HnswIndex}; +use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage}; const DIM: usize = 384; const K: usize = 10; @@ -84,6 +84,22 @@ struct Dataset { /// that appears only on clustered data points at graph connectivity. static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); +/// `--int8`: build the HNSW index over int8-quantised vectors (a quarter of +/// the memory) instead of f32, to price the recall it costs. +static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + +/// `--rerank`: re-score the candidate pool against the exact vectors before +/// taking the top K. +static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + +fn storage() -> Storage { + if INT8.load(std::sync::atomic::Ordering::Relaxed) { + Storage::Int8 + } else { + Storage::Float32 + } +} + fn make_dataset(n: usize, seed: u64) -> Dataset { let mut rng = Rng(seed); if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) { @@ -169,6 +185,11 @@ fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String { // Measurement helpers // --------------------------------------------------------------------------- +/// Exact cosine distance between unit-length vectors. +fn exact_dist(a: &[f32], b: &[f32]) -> f32 { + 1.0 - a.iter().zip(b).map(|(x, y)| x * y).sum::() +} + fn exact_top_k(vectors: &[Vec], query: &[f32], k: usize) -> Vec { // Vectors are unit length, so cosine order == dot-product order. let mut scored: Vec<(usize, f32)> = vectors @@ -270,11 +291,12 @@ fn bench_ann(n: usize, json: &mut Vec) { .collect(); let started = Instant::now(); - let index = HnswIndex::build_with_metric( + let index = HnswIndex::build_with( &data.vectors, HNSW_M, HNSW_EF_CONSTRUCTION, DistanceMetric::Cosine, + storage(), ); let build = started.elapsed(); @@ -291,7 +313,8 @@ fn bench_ann(n: usize, json: &mut Vec) { ); println!( - "\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n" + "\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}, storage = {:?}\n", + index.storage() ); println!( "build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n", @@ -302,12 +325,26 @@ fn bench_ann(n: usize, json: &mut Vec) { ); println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |"); println!("|---:|---:|---:|---:|---:|"); + // With a quantised index the distances it returns are approximate, so + // the candidates are re-scored against the exact vectors the caller + // already holds (in the agent, the embedding cache) before taking the + // top K. `--rerank` prices that: it costs one exact distance per + // candidate and is what decides whether int8 is usable. + let rerank = RERANK.load(std::sync::atomic::Ordering::Relaxed); + let pool = if rerank { K * 4 } else { K }; for ef in EF_VALUES { let mut hits = 0usize; let mut samples = Vec::with_capacity(data.queries.len()); for (q, want) in data.queries.iter().zip(&truth) { let t = Instant::now(); - let got = index.search(q, K, ef); + let mut got = index.search(q, pool, ef.max(pool)); + if rerank { + for cand in &mut got { + cand.1 = exact_dist(&data.vectors[cand.0], q); + } + got.select_nth_unstable_by(K - 1, |a, b| a.1.total_cmp(&b.1)); + got.truncate(K); + } samples.push(t.elapsed()); hits += got.iter().filter(|(id, _)| want.contains(id)).count(); } @@ -445,11 +482,12 @@ fn fusion_study(n: usize) { .map(|(i, c)| text_for(*c, i, &mut rng)) .collect(); let bm25 = BM25Index::build(&texts, &vec![0u8; n]); - let index = HnswIndex::build_with_metric( + let index = HnswIndex::build_with( &data.vectors, HNSW_M, HNSW_EF_CONSTRUCTION, DistanceMetric::Cosine, + storage(), ); let vec_pool = (K * 8).max(64); @@ -571,6 +609,14 @@ fn main() { } return; } + if args.iter().any(|a| a == "--int8") { + INT8.store(true, std::sync::atomic::Ordering::Relaxed); + println!("(int8-quantised index vectors)"); + } + if args.iter().any(|a| a == "--rerank") { + RERANK.store(true, std::sync::atomic::Ordering::Relaxed); + println!("(candidates re-scored against exact vectors)"); + } if args.iter().any(|a| a == "--uniform") { UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed); println!("(uniform random data)");