quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

perf(agent): store embeddings once, not twice

Browse Source 
MemoryCache held every embedding in two places: a `Vec<Vec<f32>>` and a
flattened copy for the batched kernels, kept in lock-step on every push,
update and compaction. A store loaded from disk therefore carried the corpus
twice, plus one heap allocation per entry.

A new `cache::Embeddings` owns just the flat `[N x dim]` buffer and indexes
into it, so `embeddings[i]` still reads as a `&[f32]` row. The batch kernels
take a `VectorSet` (implemented for both `Embeddings` and `Vec<Vec<f32>>`)
instead of `&[Vec<f32>]`, so their callers and tests are unchanged. Loading no
longer unflattens what it just read.

100k 384-dim entries, reopened from disk: 505 -> 357 MiB, 3.44x -> 2.43x the
raw vectors. Recall (1.0000 at ef=64) and query latency are unchanged.

Rows are now always exactly `dim` long, shorter ones zero-padded. The old
representation allowed ragged rows, which silently misaligned the flattened
copy — every row after a wrong-length embedding — and `update` carried a
comment about falling back to a rebuild to avoid exactly that. It is now
unrepresentable. A record saved without an embedding holds a zero row and is
told apart by its norm, which is what `total_embeddings` now counts.

Measured with a counting allocator rather than RSS: freeing a structure
returns its pages to the allocator's pool, not the OS, so an RSS reading from
inside the process showed the two representations as identical.

Breaking: MemoryCache::embeddings changes type, embeddings_flat is replaced by
flat_embeddings(), rebuild_flat() is a deprecated no-op.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-19 20:17:55 -07:00
co-authored by Claude Opus 5
parent dc0113d015
commit 2e7e0456c1
9 changed files with 429 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-agent/src/vector_search.rs
+75 -30
View File
@@ -4,6 +4,44 @@
//! `clawhdf5_accel`, with optional float16 support via the `half` crate.
//! Supports pre-computed norms for eliminating redundant norm computations.
/// A corpus of equal-length embeddings addressable by index.
///
/// Lets the batch kernels read either the cache's flat `[N x dim]` buffer or a
/// plain `Vec<Vec<f32>>` without either side owning a second copy.
pub trait VectorSet {
/// Number of embeddings.
fn count(&self) -> usize;
/// Embedding `i`; callers only index below [`VectorSet::count`].
fn row(&self, i: usize) -> &[f32];
}
impl VectorSet for [Vec<f32>] {
fn count(&self) -> usize {
self.len()
}
fn row(&self, i: usize) -> &[f32] {
&self[i]
}
}
impl VectorSet for Vec<Vec<f32>> {
fn count(&self) -> usize {
self.len()
}
fn row(&self, i: usize) -> &[f32] {
&self[i]
}
}
impl VectorSet for crate::cache::Embeddings {
fn count(&self) -> usize {
self.len()
}
fn row(&self, i: usize) -> &[f32] {
&self[i]
}
}
/// Compute cosine similarity between two f32 slices.
///
/// Returns 0.0 if either vector has zero magnitude.
@@ -22,7 +60,7 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
/// Returns `(index, score)` pairs sorted by score descending.
pub fn cosine_similarity_batch(
query: &[f32],
vectors: &[Vec<f32>],
vectors: &(impl VectorSet + ?Sized),
tombstones: &[u8],
) -> Vec<(usize, f32)> {
let query_norm = clawhdf5_accel::vector_norm(query);
@@ -30,7 +68,7 @@ pub fn cosine_similarity_batch(
return Vec::new();
}
let n = vectors.len();
let n = vectors.count();
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
// Process 4 vectors at a time where possible
@@ -42,8 +80,9 @@ pub fn cosine_similarity_batch(
if i < tombstones.len() && tombstones[i] != 0 {
continue;
}
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
let score =
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
results.push((i, score));
}
}
@@ -53,8 +92,8 @@ pub fn cosine_similarity_batch(
if i < tombstones.len() && tombstones[i] != 0 {
continue;
}
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
results.push((i, score));
}
@@ -68,7 +107,7 @@ pub fn cosine_similarity_batch(
/// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
pub fn cosine_similarity_batch_prenorm(
query: &[f32],
vectors: &[Vec<f32>],
vectors: &(impl VectorSet + ?Sized),
norms: &[f32],
tombstones: &[u8],
) -> Vec<(usize, f32)> {
@@ -77,7 +116,7 @@ pub fn cosine_similarity_batch_prenorm(
return Vec::new();
}
let n = vectors.len();
let n = vectors.count();
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
for i in 0..n {
@@ -85,7 +124,7 @@ pub fn cosine_similarity_batch_prenorm(
continue;
}
let vec_norm = norms[i];
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
results.push((i, score));
}
@@ -162,7 +201,7 @@ pub fn cosine_similarity_f16(
#[cfg(feature = "parallel")]
pub fn parallel_cosine_batch(
query: &[f32],
vectors: &[Vec<f32>],
vectors: &(impl VectorSet + Sync + ?Sized),
tombstones: &[u8],
k: usize,
) -> Vec<(usize, f32)> {
@@ -174,24 +213,27 @@ pub fn parallel_cosine_batch(
}
let num_cores = rayon::current_num_threads().max(1);
let chunk_size = vectors.len().div_ceil(num_cores);
let chunk_size = vectors.count().div_ceil(num_cores);
if chunk_size == 0 {
return Vec::new();
}
let mut all_results: Vec<(usize, f32)> = vectors
.par_chunks(chunk_size)
.enumerate()
.flat_map(|(chunk_idx, chunk)| {
// Chunk over index ranges: the corpus may be one flat buffer rather than
// a slice of rows, so there is nothing to `par_chunks` over.
let n = vectors.count();
let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size))
.into_par_iter()
.flat_map(|chunk_idx| {
let base = chunk_idx * chunk_size;
let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
for (j, vec) in chunk.iter().enumerate() {
let i = base + j;
let end = (base + chunk_size).min(n);
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
for i in base..end {
if i < tombstones.len() && tombstones[i] != 0 {
continue;
}
let vec_norm = clawhdf5_accel::vector_norm(vec);
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm);
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
let score =
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
local.push((i, score));
}
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
@@ -209,7 +251,7 @@ pub fn parallel_cosine_batch(
#[cfg(feature = "parallel")]
pub fn parallel_cosine_batch_prenorm(
query: &[f32],
vectors: &[Vec<f32>],
vectors: &(impl VectorSet + Sync + ?Sized),
norms: &[f32],
tombstones: &[u8],
k: usize,
@@ -222,23 +264,26 @@ pub fn parallel_cosine_batch_prenorm(
}
let num_cores = rayon::current_num_threads().max(1);
let chunk_size = vectors.len().div_ceil(num_cores);
let chunk_size = vectors.count().div_ceil(num_cores);
if chunk_size == 0 {
return Vec::new();
}
let mut all_results: Vec<(usize, f32)> = vectors
.par_chunks(chunk_size)
.enumerate()
.flat_map(|(chunk_idx, chunk)| {
// Chunk over index ranges: the corpus may be one flat buffer rather than
// a slice of rows, so there is nothing to `par_chunks` over.
let n = vectors.count();
let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size))
.into_par_iter()
.flat_map(|chunk_idx| {
let base = chunk_idx * chunk_size;
let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
for (j, vec) in chunk.iter().enumerate() {
let i = base + j;
let end = (base + chunk_size).min(n);
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
for i in base..end {
if i < tombstones.len() && tombstones[i] != 0 {
continue;
}
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, norms[i]);
let score =
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), norms[i]);
local.push((i, score));
}
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));