perf(agent): store embeddings once, not twice
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:
@@ -565,7 +565,9 @@ impl HDF5Memory {
|
||||
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
||||
return None;
|
||||
}
|
||||
let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec();
|
||||
let vectors: Vec<Vec<f32>> = (0..n_checkpoint)
|
||||
.map(|i| cache.embeddings.get(i).map(<[f32]>::to_vec))
|
||||
.collect::<Option<_>>()?;
|
||||
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
|
||||
if index.dimension() != cache.embedding_dim {
|
||||
return None;
|
||||
@@ -573,7 +575,7 @@ impl HDF5Memory {
|
||||
// Records appended since (replayed from the WAL) join incrementally.
|
||||
for id in n_checkpoint..cache.embeddings.len() {
|
||||
if cache.embeddings[id].len() != index.dimension()
|
||||
|| index.insert(cache.embeddings[id].clone()) != id
|
||||
|| index.insert(cache.embeddings[id].to_vec()) != id
|
||||
{
|
||||
return None;
|
||||
}
|
||||
@@ -816,8 +818,11 @@ impl HDF5Memory {
|
||||
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
||||
return None;
|
||||
}
|
||||
// The index owns its vectors, so it needs rows rather than the cache's
|
||||
// flat buffer. This copy is the index's own; the cache keeps one.
|
||||
let rows: Vec<Vec<f32>> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect();
|
||||
let mut index = HnswIndex::build_with_metric(
|
||||
&self.cache.embeddings,
|
||||
&rows,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
@@ -846,7 +851,7 @@ impl HDF5Memory {
|
||||
let dim = index.dimension();
|
||||
let appended = (self.hnsw_synced_len..n).all(|id| {
|
||||
self.cache.embeddings[id].len() == dim
|
||||
&& index.insert(self.cache.embeddings[id].clone()) == id
|
||||
&& index.insert(self.cache.embeddings[id].to_vec()) == id
|
||||
});
|
||||
if appended {
|
||||
for id in self.hnsw_synced_len..n {
|
||||
@@ -877,7 +882,7 @@ impl HDF5Memory {
|
||||
let emb_len = self.cache.embeddings[idx].len();
|
||||
match self.hnsw.as_mut() {
|
||||
Some(index) if emb_len == index.dimension() => {
|
||||
let id = index.insert(self.cache.embeddings[idx].clone());
|
||||
let id = index.insert(self.cache.embeddings[idx].to_vec());
|
||||
if id == idx {
|
||||
self.hnsw_synced_len = self.cache.embeddings.len();
|
||||
} else {
|
||||
|
||||
Reference in New Issue
Block a user