perf(ann): unit-vector dot product and a reusable visited set
- Cosine distance was 1 - dot/(|a||b|), re-deriving both norms on every evaluation in the innermost loop of build and search. The index now stores unit vectors (prepared at build, insert, graph load and HDF5 load; the query once per search) and uses 1 - dot. Zero vectors stay zero, giving distance 1 as before. Returned distances are unchanged. - search_layer allocated a HashSet of visited nodes per call. It is now an epoch-stamped u32 array in thread-local scratch, reused across calls, so search(&self) stays shareable between threads. - clawhdf5-accel caches the detected SIMD backend in a OnceLock. Recall is identical. Build 2.75 -> 1.89 s (10K), ~38 -> 21 s (100K); QPS at ef=64 22.7K -> 39K (10K), 10.4K -> 14K (100K). Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
09480747aa
commit
f15bf2eb22
@@ -176,6 +176,51 @@ index incrementally.
|
||||
| 10000 | 120 | 2916 | 33.1 | 14.0 | 15.4 | 2.15 | 3.30 | 421.2 |
|
||||
| 100000 | 1591 | 40515 | 747.3 | 324.7 | 158.9 | 23.07 | 30.42 | 41.3 |
|
||||
|
||||
### After: unit-vector dot product, reusable visited set
|
||||
|
||||
Cosine distance recomputed both vector norms on every evaluation; the index now
|
||||
stores unit vectors and uses a plain dot product. The per-call `HashSet` of
|
||||
visited nodes became a reusable epoch-stamped array. Recall is unchanged.
|
||||
Build: **2.75 -> 1.89 s** (10K), **~38 -> 21 s** (100K). QPS at `ef = 64`:
|
||||
**22.7K -> 39K** (10K), **10.4K -> 14K** (100K).
|
||||
|
||||
### HNSW, N = 1000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 113.4 ms (8821 vectors/s) · exact scan: 4375 QPS, p50 225 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9990 | 144379 | 7 | 15 |
|
||||
| 32 | 1.0000 | 110654 | 9 | 17 |
|
||||
| 64 | 1.0000 | 80446 | 12 | 25 |
|
||||
| 128 | 1.0000 | 38220 | 26 | 36 |
|
||||
| 256 | 1.0000 | 20041 | 50 | 62 |
|
||||
|
||||
### HNSW, N = 10000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 1519.4 ms (6581 vectors/s) · exact scan: 422 QPS, p50 2368 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9975 | 54608 | 15 | 45 |
|
||||
| 32 | 1.0000 | 66009 | 14 | 24 |
|
||||
| 64 | 1.0000 | 49854 | 19 | 31 |
|
||||
| 128 | 1.0000 | 22403 | 45 | 57 |
|
||||
| 256 | 1.0000 | 10096 | 100 | 120 |
|
||||
|
||||
### HNSW, N = 100000, dim = 384, M = 16, ef_construction = 64
|
||||
|
||||
build: 21084.6 ms (4743 vectors/s) · exact scan: 39 QPS, p50 24739 µs
|
||||
|
||||
| ef | recall@10 | QPS | p50 µs | p99 µs |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 16 | 0.9235 | 15139 | 61 | 154 |
|
||||
| 32 | 0.9675 | 18181 | 53 | 121 |
|
||||
| 64 | 0.9840 | 13980 | 70 | 139 |
|
||||
| 128 | 0.9990 | 10959 | 86 | 174 |
|
||||
| 256 | 0.9990 | 3731 | 254 | 697 |
|
||||
|
||||
|
||||
## Vector Search Latency
|
||||
|
||||
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
|
||||
|
||||
@@ -32,6 +32,15 @@
|
||||
replayed from the WAL join the loaded index incrementally; a replayed update
|
||||
or delete invalidates it. `snapshot()` copies it. Batch saves no longer force
|
||||
a full index rebuild.
|
||||
- `clawhdf5-ann`: faster HNSW build and search with identical recall. The
|
||||
cosine metric stores unit vectors and compares them with a plain dot product
|
||||
(it re-derived both norms on every distance evaluation), and the per-call
|
||||
`HashSet` of visited nodes is a reusable epoch-stamped array. Build 2.75 ->
|
||||
1.89 s at 10K and ~38 -> 21 s at 100K; QPS at `ef = 64` 22.7K -> 39K at 10K.
|
||||
Distances returned by `search` are unchanged (1 - cosine). Indexes loaded
|
||||
from older HDF5 files are normalised on load.
|
||||
- `clawhdf5-accel`: the SIMD backend is detected once per process instead of
|
||||
on every kernel call.
|
||||
- `clawhdf5-ann`: `HnswIndex::graph_to_bytes` / `from_graph_bytes` — graph-only
|
||||
serialization (checksummed, every neighbour id and level validated on load).
|
||||
- `clawhdf5-agent`: BM25 results are deterministic (ties break by record id),
|
||||
|
||||
@@ -61,8 +61,14 @@ pub enum Backend {
|
||||
Scalar,
|
||||
}
|
||||
|
||||
/// Detect the best available SIMD backend at runtime.
|
||||
/// The best available SIMD backend, detected once per process. Every kernel
|
||||
/// dispatches through this, so it sits in the innermost loop of every search.
|
||||
pub fn detect_backend() -> Backend {
|
||||
static BACKEND: std::sync::OnceLock<Backend> = std::sync::OnceLock::new();
|
||||
*BACKEND.get_or_init(detect_backend_uncached)
|
||||
}
|
||||
|
||||
fn detect_backend_uncached() -> Backend {
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
{
|
||||
return Backend::Neon; // Always available on aarch64
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! HNSW index implementation with HDF5 serialization.
|
||||
|
||||
use std::collections::{BinaryHeap, HashSet};
|
||||
use std::collections::BinaryHeap;
|
||||
|
||||
use clawhdf5_format::attribute::extract_attributes_full;
|
||||
use clawhdf5_format::data_layout::DataLayout;
|
||||
@@ -51,10 +51,30 @@ impl DistanceMetric {
|
||||
fn compute_distance(a: &[f32], b: &[f32], metric: DistanceMetric) -> f32 {
|
||||
match metric {
|
||||
DistanceMetric::L2 => clawhdf5_accel::l2_distance(a, b),
|
||||
DistanceMetric::Cosine => 1.0 - clawhdf5_accel::cosine_similarity(a, b),
|
||||
// Both sides are unit length (see `prepare`), so cosine similarity is
|
||||
// the plain dot product. Computing it as dot / (|a| * |b|) re-derived
|
||||
// both norms on every call — three reductions instead of one, in the
|
||||
// innermost loop of both build and search.
|
||||
DistanceMetric::Cosine => 1.0 - clawhdf5_accel::dot_product(a, b),
|
||||
}
|
||||
}
|
||||
|
||||
/// Put a vector in the form the index stores and compares: unit length for the
|
||||
/// cosine metric, unchanged for L2. A zero vector stays zero, giving distance 1
|
||||
/// to everything — what the cosine kernel reports for a degenerate input.
|
||||
fn prepare(mut v: Vec<f32>, metric: DistanceMetric) -> Vec<f32> {
|
||||
if metric == DistanceMetric::Cosine {
|
||||
let norm = clawhdf5_accel::vector_norm(&v);
|
||||
if norm > f32::EPSILON {
|
||||
let inv = 1.0 / norm;
|
||||
v.iter_mut().for_each(|x| *x *= inv);
|
||||
} else {
|
||||
v.iter_mut().for_each(|x| *x = 0.0);
|
||||
}
|
||||
}
|
||||
v
|
||||
}
|
||||
|
||||
/// Assign a random level to a new node based on the HNSW probability distribution.
|
||||
///
|
||||
/// Uses a deterministic approach based on the node index for reproducibility.
|
||||
@@ -204,6 +224,8 @@ impl HnswIndex {
|
||||
|
||||
let m_max0 = m * 2;
|
||||
let n = vectors.len();
|
||||
let prepared: Vec<Vec<f32>> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect();
|
||||
let vectors: &[Vec<f32>] = &prepared;
|
||||
|
||||
// Assign levels to all nodes
|
||||
let mut node_levels = Vec::with_capacity(n);
|
||||
@@ -288,7 +310,7 @@ impl HnswIndex {
|
||||
}
|
||||
|
||||
Self {
|
||||
vectors: vectors.to_vec(),
|
||||
vectors: prepared,
|
||||
graph,
|
||||
deleted: vec![false; n],
|
||||
entry_point,
|
||||
@@ -327,6 +349,7 @@ impl HnswIndex {
|
||||
/// # Panics
|
||||
/// Panics if `vector`'s dimension does not match the existing vectors.
|
||||
pub fn insert(&mut self, vector: Vec<f32>) -> usize {
|
||||
let vector = prepare(vector, self.metric);
|
||||
let id = self.vectors.len();
|
||||
|
||||
// Seed an empty index.
|
||||
@@ -483,6 +506,8 @@ impl HnswIndex {
|
||||
"query dimension mismatch"
|
||||
);
|
||||
let ef = ef.max(k);
|
||||
let prepared_query = prepare(query.to_vec(), self.metric);
|
||||
let query = prepared_query.as_slice();
|
||||
|
||||
let mut ep = self.entry_point;
|
||||
let top_layer = self.graph.len().saturating_sub(1);
|
||||
@@ -633,7 +658,9 @@ impl HnswIndex {
|
||||
actual: flat_vectors.len(),
|
||||
});
|
||||
}
|
||||
vectors.push(flat_vectors[start..end].to_vec());
|
||||
// Files written before vectors were stored unit-length hold the
|
||||
// raw ones; preparing is idempotent, so this handles both.
|
||||
vectors.push(prepare(flat_vectors[start..end].to_vec(), metric));
|
||||
}
|
||||
|
||||
// Read graph layers
|
||||
@@ -835,7 +862,7 @@ impl HnswIndex {
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
vectors,
|
||||
vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(),
|
||||
graph,
|
||||
deleted,
|
||||
entry_point,
|
||||
@@ -937,9 +964,62 @@ fn search_layer(
|
||||
distance: ep_dist,
|
||||
});
|
||||
|
||||
let mut visited = HashSet::new();
|
||||
VISITED.with_borrow_mut(|visited| {
|
||||
visited.begin(vectors.len());
|
||||
visited.insert(ep);
|
||||
search_layer_visit(
|
||||
vectors, layer, query, ef, metric, visited, candidates, results,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Which nodes a layer search has already seen. A `HashSet` allocated per call
|
||||
/// was the hottest non-arithmetic cost in both build and query; this is one
|
||||
/// `u32` stamp per node, reused across calls: a node is visited iff its stamp
|
||||
/// equals the current epoch, so "clearing" is just bumping the epoch.
|
||||
#[derive(Default)]
|
||||
struct Visited {
|
||||
stamps: Vec<u32>,
|
||||
epoch: u32,
|
||||
}
|
||||
|
||||
impl Visited {
|
||||
fn begin(&mut self, n: usize) {
|
||||
if self.stamps.len() < n {
|
||||
self.stamps.resize(n, 0);
|
||||
}
|
||||
self.epoch = self.epoch.wrapping_add(1);
|
||||
if self.epoch == 0 {
|
||||
// Wrapped: stale stamps could collide with the new epoch.
|
||||
self.stamps.iter_mut().for_each(|s| *s = 0);
|
||||
self.epoch = 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark `id` visited; `true` if it was not already.
|
||||
fn insert(&mut self, id: usize) -> bool {
|
||||
let seen = self.stamps[id] == self.epoch;
|
||||
self.stamps[id] = self.epoch;
|
||||
!seen
|
||||
}
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
/// Per-thread scratch, so `search(&self)` stays shareable across threads.
|
||||
static VISITED: std::cell::RefCell<Visited> = std::cell::RefCell::new(Visited::default());
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn search_layer_visit(
|
||||
vectors: &[Vec<f32>],
|
||||
layer: &[Vec<usize>],
|
||||
query: &[f32],
|
||||
ef: usize,
|
||||
metric: DistanceMetric,
|
||||
visited: &mut Visited,
|
||||
mut candidates: BinaryHeap<Candidate>,
|
||||
mut results: BinaryHeap<FarCandidate>,
|
||||
) -> Vec<Candidate> {
|
||||
while let Some(closest) = candidates.pop() {
|
||||
let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance);
|
||||
if closest.distance > furthest_dist && results.len() >= ef {
|
||||
@@ -947,10 +1027,9 @@ fn search_layer(
|
||||
}
|
||||
|
||||
for &neighbor in &layer[closest.id] {
|
||||
if visited.contains(&neighbor) {
|
||||
if !visited.insert(neighbor) {
|
||||
continue;
|
||||
}
|
||||
visited.insert(neighbor);
|
||||
|
||||
let d = compute_distance(query, &vectors[neighbor], metric);
|
||||
let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance);
|
||||
@@ -1236,6 +1315,7 @@ fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String,
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::HashSet;
|
||||
|
||||
/// Tight, well-separated clusters — the shape real embeddings have, and
|
||||
/// the case plain closest-M neighbour selection fails on: each cluster
|
||||
|
||||
Reference in New Issue
Block a user