feat(agent): optional int8 vector index, re-scored against exact embeddings
`MemoryConfig::quantized_index` stores the HNSW index's own copy of the embeddings as i8 rather than f32. At 100k x 384 that takes the index from 266 to 123 MiB and the whole reopened store from 399 to 256 MiB — 2.72x to 1.74x the raw vectors, the largest remaining item in the footprint. Quantised distances are approximate and `ef` cannot compensate, because the loss is in the distances rather than in the graph: recall@10 tops out at 0.967 against f32's 0.9995 and does not move between ef=128 and ef=256. The store already holds the exact embeddings, though, so when the index is quantised the query path re-scores the candidate pool against them before fusion. That restores recall (0.9940 vs 0.9945 at ef=64) and costs about 13% of QPS. Off by default: it trades query speed for memory and which side is worth more depends on the deployment. The flag is persisted in `/meta`, so a reopened store does not silently revert to four times the index memory, and the sidecar graph is rehydrated into the configured storage. Also on the CLI as `create --quantized-index`. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -62,7 +62,7 @@ use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
|
||||
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
||||
@@ -139,6 +139,17 @@ pub struct MemoryConfig {
|
||||
pub created_at: String,
|
||||
pub wal_enabled: bool,
|
||||
pub wal_max_entries: usize,
|
||||
/// Store the vector index's own copy of the embeddings as int8 rather than
|
||||
/// f32, a quarter of the memory.
|
||||
///
|
||||
/// The index's copy is the single largest part of a loaded store's
|
||||
/// footprint. Quantised distances are approximate, so the candidate pool
|
||||
/// is re-scored against the cache's exact embeddings before fusion, which
|
||||
/// restores recall; what it costs is throughput — roughly 13% of queries
|
||||
/// per second and 16% of build time at 100K x 384. See `BENCHMARKS.md`.
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub quantized_index: bool,
|
||||
}
|
||||
|
||||
impl MemoryConfig {
|
||||
@@ -160,6 +171,7 @@ impl MemoryConfig {
|
||||
created_at,
|
||||
wal_enabled: true,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -444,7 +456,17 @@ impl HDF5Memory {
|
||||
|
||||
#[cfg(feature = "hnsw")]
|
||||
let loaded_index = if replay_only_appended {
|
||||
Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint)
|
||||
Self::load_vector_index(
|
||||
path,
|
||||
checkpoint.ann_generation,
|
||||
&cache,
|
||||
n_checkpoint,
|
||||
if config.quantized_index {
|
||||
Storage::Int8
|
||||
} else {
|
||||
Storage::Float32
|
||||
},
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -558,6 +580,7 @@ impl HDF5Memory {
|
||||
generation: Option<u64>,
|
||||
cache: &MemoryCache,
|
||||
n_checkpoint: usize,
|
||||
storage: Storage,
|
||||
) -> Option<HnswIndex> {
|
||||
let generation = generation?;
|
||||
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
|
||||
@@ -568,7 +591,7 @@ impl HDF5Memory {
|
||||
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()?;
|
||||
let mut index = HnswIndex::from_graph_bytes_with(graph, vectors, storage).ok()?;
|
||||
if index.dimension() != cache.embedding_dim {
|
||||
return None;
|
||||
}
|
||||
@@ -803,6 +826,16 @@ impl HDF5Memory {
|
||||
// the index length drifts from the cache length (covering any mutation path
|
||||
// that doesn't call a hook, e.g. consolidation pushes).
|
||||
|
||||
/// How the index should store its copy of the vectors, per the config.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn index_storage(&self) -> Storage {
|
||||
if self.config.quantized_index {
|
||||
Storage::Int8
|
||||
} else {
|
||||
Storage::Float32
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an HNSW index over the entire cache, re-applying tombstones as
|
||||
/// soft-deletions so node ids stay aligned with cache indices.
|
||||
///
|
||||
@@ -821,11 +854,12 @@ impl HDF5Memory {
|
||||
// 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(
|
||||
let mut index = HnswIndex::build_with(
|
||||
&rows,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistanceMetric::Cosine,
|
||||
self.index_storage(),
|
||||
);
|
||||
for (i, &t) in self.cache.tombstones.iter().enumerate() {
|
||||
if t != 0 {
|
||||
|
||||
Reference in New Issue
Block a user