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:
@@ -153,7 +153,7 @@ fn build_memory_group(
|
||||
{
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_f32_data(&flat)
|
||||
.with_f32_data(flat)
|
||||
.with_shape(&[n, d]);
|
||||
|
||||
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
||||
@@ -563,12 +563,7 @@ fn load_memory_group(
|
||||
.collect(),
|
||||
};
|
||||
|
||||
// Unflatten embeddings
|
||||
let embeddings: Vec<Vec<f32>> = flat_embeddings
|
||||
.chunks(embedding_dim)
|
||||
.map(|c| c.to_vec())
|
||||
.collect();
|
||||
|
||||
// No unflattening: the cache stores the buffer as it is on disk.
|
||||
// Read activation_weights if present, default to vec![1.0; N] for backward compat
|
||||
let activation_weights = match read_f32_dataset(&group, "activation_weights") {
|
||||
Ok(w) if w.len() == n => w,
|
||||
@@ -576,7 +571,7 @@ fn load_memory_group(
|
||||
};
|
||||
|
||||
cache.chunks = chunks;
|
||||
cache.embeddings = embeddings;
|
||||
cache.embeddings.set_flat(embedding_dim, flat_embeddings);
|
||||
cache.source_channels = source_channels;
|
||||
cache.timestamps = timestamps;
|
||||
cache.session_ids = session_ids;
|
||||
@@ -584,7 +579,6 @@ fn load_memory_group(
|
||||
cache.tombstones = tombstones;
|
||||
cache.norms = norms;
|
||||
cache.activation_weights = activation_weights;
|
||||
cache.rebuild_flat();
|
||||
|
||||
Ok(cache)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user