crates/clawhdf5-agent/src/cache.rs

@@ -7,6 +7,11 @@ use crate::vector_search;
 pub struct MemoryCache {
     pub chunks: Vec<String>,
     pub embeddings: Vec<Vec<f32>>,
+    /// `embeddings` flattened into one contiguous `[N × embedding_dim]`
+    /// buffer, maintained incrementally alongside `embeddings` (push/update/
+    /// compact) so BLAS/Accelerate batch search can read it directly instead
+    /// of re-flattening the whole corpus on every query.
+    pub embeddings_flat: Vec<f32>,
     pub source_channels: Vec<String>,
     pub timestamps: Vec<f64>,
     pub session_ids: Vec<String>,
@@ -24,6 +29,7 @@ impl MemoryCache {
         Self {
             chunks: Vec::new(),
             embeddings: Vec::new(),
+            embeddings_flat: Vec::new(),
             source_channels: Vec::new(),
             timestamps: Vec::new(),
             session_ids: Vec::new(),
@@ -35,6 +41,17 @@ impl MemoryCache {
         }
     }
 
+    /// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that
+    /// populate `embeddings` directly (bulk loads) must call this afterward.
+    pub fn rebuild_flat(&mut self) {
+        self.embeddings_flat.clear();
+        self.embeddings_flat
+            .reserve(self.embeddings.len() * self.embedding_dim);
+        for emb in &self.embeddings {
+            self.embeddings_flat.extend_from_slice(emb);
+        }
+    }
+
     /// Total number of entries (including tombstoned).
     pub fn len(&self) -> usize {
         self.chunks.len()
@@ -62,6 +79,7 @@ impl MemoryCache {
         let idx = self.chunks.len();
         let norm = vector_search::compute_norm(&embedding);
         self.chunks.push(chunk);
+        self.embeddings_flat.extend_from_slice(&embedding);
         self.embeddings.push(embedding);
         self.source_channels.push(source_channel);
         self.timestamps.push(timestamp);
@@ -100,7 +118,20 @@ impl MemoryCache {
         if idx < self.chunks.len() {
             let norm = vector_search::compute_norm(&embedding);
             self.chunks[idx] = chunk;
+            let dim = self.embedding_dim;
+            let flat_start = idx * dim;
+            let matches_dim =
+                embedding.len() == dim && flat_start + dim <= self.embeddings_flat.len();
             self.embeddings[idx] = embedding;
+            if matches_dim {
+                self.embeddings_flat[flat_start..flat_start + dim]
+                    .copy_from_slice(&self.embeddings[idx]);
+            } else {
+                // Embedding length doesn't match embedding_dim (shouldn't
+                // happen in practice) — fall back to a full rebuild rather
+                // than leave embeddings_flat misaligned with embeddings.
+                self.rebuild_flat();
+            }
             self.source_channels[idx] = source_channel;
             self.timestamps[idx] = timestamp;
             self.session_ids[idx] = session_id;
@@ -173,16 +204,125 @@ impl MemoryCache {
         self.tombstones = new_tombstones;
         self.norms = new_norms;
         self.activation_weights = new_activation_weights;
+        self.rebuild_flat();
 
         (removed, index_map)
     }
 
     /// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
+    /// `embeddings_flat` is already maintained incrementally, so this just
+    /// clones it — kept as a method for callers that want an owned copy.
     pub fn flat_embeddings(&self) -> Vec<f32> {
-        let mut flat = Vec::with_capacity(self.embeddings.len() * self.embedding_dim);
+        self.embeddings_flat.clone()
-        for emb in &self.embeddings {
-            flat.extend_from_slice(emb);
     }
-        flat
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
+    fn assert_flat_in_sync(cache: &MemoryCache) {
+        let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
+        assert_eq!(cache.embeddings_flat, expected);
+    }
+
+    #[test]
+    fn push_keeps_flat_buffer_in_sync() {
+        let mut cache = MemoryCache::new(3);
+        cache.push(
+            "a".into(),
+            vec![1.0, 2.0, 3.0],
+            "chan".into(),
+            0.0,
+            "s1".into(),
+            String::new(),
+        );
+        cache.push(
+            "b".into(),
+            vec![4.0, 5.0, 6.0],
+            "chan".into(),
+            1.0,
+            "s1".into(),
+            String::new(),
+        );
+        assert_flat_in_sync(&cache);
+        assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
+    }
+
+    #[test]
+    fn update_keeps_flat_buffer_in_sync() {
+        let mut cache = MemoryCache::new(3);
+        cache.push(
+            "a".into(),
+            vec![1.0, 2.0, 3.0],
+            "chan".into(),
+            0.0,
+            "s1".into(),
+            String::new(),
+        );
+        cache.push(
+            "b".into(),
+            vec![4.0, 5.0, 6.0],
+            "chan".into(),
+            1.0,
+            "s1".into(),
+            String::new(),
+        );
+        cache.update(
+            0,
+            "a2".into(),
+            vec![7.0, 8.0, 9.0],
+            "chan".into(),
+            2.0,
+            "s1".into(),
+        );
+        assert_flat_in_sync(&cache);
+        assert_eq!(
+            cache.embeddings_flat,
+            vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
+            "update must overwrite the correct flat slice, not just append"
+        );
+    }
+
+    #[test]
+    fn compact_keeps_flat_buffer_in_sync() {
+        let mut cache = MemoryCache::new(2);
+        cache.push(
+            "a".into(),
+            vec![1.0, 1.0],
+            "chan".into(),
+            0.0,
+            "s1".into(),
+            String::new(),
+        );
+        cache.push(
+            "b".into(),
+            vec![2.0, 2.0],
+            "chan".into(),
+            1.0,
+            "s1".into(),
+            String::new(),
+        );
+        cache.push(
+            "c".into(),
+            vec![3.0, 3.0],
+            "chan".into(),
+            2.0,
+            "s1".into(),
+            String::new(),
+        );
+        cache.mark_deleted(1);
+        cache.compact();
+        assert_flat_in_sync(&cache);
+        assert_eq!(cache.embeddings_flat, vec![1.0, 1.0, 3.0, 3.0]);
+    }
+
+    #[test]
+    fn rebuild_flat_matches_manual_flatten() {
+        let mut cache = MemoryCache::new(2);
+        cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
+        cache.rebuild_flat();
+        assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]);
     }
 }

crates/clawhdf5-agent/src/schema.rs

@@ -427,6 +427,7 @@ fn load_memory_group(
     cache.tombstones = tombstones;
     cache.norms = norms;
     cache.activation_weights = activation_weights;
+    cache.rebuild_flat();
 
     Ok(cache)
 }

crates/clawhdf5-agent/src/strategy.rs

@@ -167,10 +167,17 @@ pub fn auto_select_strategy(num_vectors: usize, hw: &HardwareCapabilities) -> Se
 /// This dispatches to the appropriate search implementation based on the
 /// selected strategy. For IVF-PQ, an index must be provided externally
 /// (this function uses brute-force fallback if no IVF-PQ index is available).
+///
+/// `vectors_flat` is `vectors` flattened into one contiguous `[N × dim]`
+/// row-major buffer (e.g. `MemoryCache::embeddings_flat`, maintained
+/// incrementally alongside `vectors`). It's only consulted by the
+/// `Blas`/`Accelerate` strategies, which otherwise re-flatten the whole
+/// corpus on every call — passing the already-flat buffer skips that copy.
 #[allow(clippy::too_many_arguments)]
 pub fn search_with_metrics(
     query: &[f32],
     vectors: &[Vec<f32>],
+    vectors_flat: &[f32],
     norms: &[f32],
     tombstones: &[u8],
     k: usize,
@@ -178,6 +185,10 @@ pub fn search_with_metrics(
     #[cfg(feature = "gpu")] gpu_backend: Option<&crate::gpu_search::GpuSearchBackend>,
     #[cfg(not(feature = "gpu"))] _gpu_backend: Option<&()>,
 ) -> (Vec<(usize, f32)>, SearchMetrics) {
+    // Only read by the Blas/Accelerate arms below, which are themselves
+    // feature-gated — reference it unconditionally so a build with neither
+    // feature enabled doesn't warn about an unused parameter.
+    let _ = vectors_flat;
     let start = Instant::now();
     let active_count = tombstones.iter().filter(|&&t| t == 0).count();
 
@@ -197,7 +208,14 @@ pub fn search_with_metrics(
             gpu_active = false;
             #[cfg(feature = "fast-math")]
             {
-                crate::blas_search::blas_cosine_batch(query, vectors, norms, tombstones, k)
+                crate::blas_search::blas_cosine_batch_flat(
+                    query,
+                    vectors_flat,
+                    norms,
+                    tombstones,
+                    query.len(),
+                    k,
+                )
             }
             #[cfg(not(feature = "fast-math"))]
             {
@@ -211,8 +229,13 @@ pub fn search_with_metrics(
             gpu_active = false;
             #[cfg(any(feature = "accelerate", feature = "openblas"))]
             {
-                crate::accelerate_search::accelerate_cosine_batch_vecs(
+                crate::accelerate_search::accelerate_cosine_batch(
-                    query, vectors, norms, tombstones, k,
+                    query,
+                    vectors_flat,
+                    norms,
+                    tombstones,
+                    query.len(),
+                    k,
                 )
             }
             #[cfg(not(any(feature = "accelerate", feature = "openblas")))]
@@ -325,6 +348,10 @@ mod tests {
         (0..n).map(|_| (0..dim).map(|_| next()).collect()).collect()
     }
 
+    fn flatten(vectors: &[Vec<f32>]) -> Vec<f32> {
+        vectors.iter().flatten().copied().collect()
+    }
+
     // --- auto_select_strategy tests ---
 
     #[test]
@@ -490,6 +517,7 @@ mod tests {
         let (results, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             5,
@@ -520,6 +548,7 @@ mod tests {
         let (results, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             10,
@@ -545,6 +574,7 @@ mod tests {
         let (_, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             10,
@@ -570,6 +600,7 @@ mod tests {
         let (results, _) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             10,
@@ -603,6 +634,7 @@ mod tests {
         let (results, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             100,
@@ -647,6 +679,7 @@ mod tests {
         let (_, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             5,
@@ -718,6 +751,7 @@ mod tests {
         let (results, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             10,
@@ -744,6 +778,7 @@ mod tests {
         let (results, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             10,
@@ -822,6 +857,7 @@ mod tests {
         let (results, metrics) = search_with_metrics(
             &query,
             &vectors,
+            &flatten(&vectors),
             &norms,
             &tombstones,
             10,