diff --git a/Cargo.lock b/Cargo.lock index 03d9322bc8..051a742a93 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4938,6 +4938,24 @@ dependencies = [ "tracing", ] +[[package]] +name = "lattice-embed" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a8471670d8eb3dc5b52b7c977f1be449a4d39404ebd47bd084a1e922fa6002e" +dependencies = [ + "async-trait", + "blake3", + "chrono", + "lru 0.16.4", + "parking_lot 0.12.5", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "lattice-inference" version = "0.6.1" @@ -9203,7 +9221,7 @@ dependencies = [ "dashmap 6.2.1", "hf-hub", "hnsw_rs", - "lattice-embed", + "lattice-embed 0.6.1", "memmap2", "mockall", "ndarray 0.16.1", @@ -10698,6 +10716,7 @@ dependencies = [ "criterion 0.5.1", "crossbeam", "dashmap 6.2.1", + "lattice-embed 0.7.1", "memmap2", "ndarray 0.16.1", "once_cell", diff --git a/crates/ruvector-tiny-dancer-core/Cargo.toml b/crates/ruvector-tiny-dancer-core/Cargo.toml index 1dff762082..b6fff73d2a 100644 --- a/crates/ruvector-tiny-dancer-core/Cargo.toml +++ b/crates/ruvector-tiny-dancer-core/Cargo.toml @@ -11,6 +11,14 @@ description = "Production-grade AI agent routing system with FastGRNN neural inf [lib] crate-type = ["lib", "staticlib"] +[features] +# The cosine kernel behind candidate scoring. Exactly one backend compiles; +# `lattice-simd` wins if both are on. `--no-default-features` leaves the scalar +# path, which is also the reference the parity test checks the others against. +default = ["simd-simsimd"] +simd-simsimd = ["dep:simsimd"] +lattice-simd = ["dep:lattice-embed"] + [dependencies] # Workspace dependencies redb = { workspace = true } @@ -20,7 +28,15 @@ crossbeam = { workspace = true } parking_lot = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -simsimd = { workspace = true } +simsimd = { workspace = true, optional = true } +# Pure-Rust SIMD kernels with runtime dispatch, and the only one of the three +# backends that compiles for wasm32. `default-features = false` keeps this to +# the kernels: the model, tokenizer and download stack stay out of the graph. +# +# NOTE: lattice-embed requires Rust >= 1.93. Cargo cannot express a per-feature +# `rust-version`, so enabling `lattice-simd` raises the effective MSRV for +# whoever turns it on. The default build is unaffected. +lattice-embed = { version = "0.7.1", optional = true, default-features = false } thiserror = { workspace = true } anyhow = { workspace = true } tracing = { workspace = true } diff --git a/crates/ruvector-tiny-dancer-core/README.md b/crates/ruvector-tiny-dancer-core/README.md index 01ea2ee6cb..de90a0f894 100644 --- a/crates/ruvector-tiny-dancer-core/README.md +++ b/crates/ruvector-tiny-dancer-core/README.md @@ -322,10 +322,32 @@ let scores = model.forward_batch(&inputs)?; ### SIMD Acceleration -Feature extraction uses `simsimd` for hardware-accelerated similarity: -- Cosine similarity: **144ns** (384-dim vectors) +Feature extraction's cosine similarity kernel is selected by Cargo feature: + +| Feature | Backend | Notes | +|---------|---------|-------| +| `simd-simsimd` (default) | `simsimd` | Unchanged default behavior; the 144ns benchmark below is this path. | +| `lattice-simd` | `lattice-embed` | Opt-in. Wins at runtime if both features are enabled. Requires Rust >= 1.93; the default build stays on the workspace MSRV. | +| *(no features)* | scalar fallback | `--no-default-features` with neither backend enabled. | + +- Cosine similarity: **144ns** (384-dim vectors, default `simd-simsimd` backend) - Batch processing: **Linear scaling** with candidate count +```sh +# Default: simd-simsimd +cargo build -p ruvector-tiny-dancer-core + +# Opt into lattice-simd (still compiles simsimd into the dependency graph, +# since features are additive; lattice wins the runtime selection) +cargo build -p ruvector-tiny-dancer-core --features lattice-simd + +# Lattice only, with simsimd excluded from the dependency graph entirely +cargo build -p ruvector-tiny-dancer-core --no-default-features --features lattice-simd + +# Scalar fallback, no SIMD backend compiled in +cargo build -p ruvector-tiny-dancer-core --no-default-features +``` + ### Zero-Copy Operations - Memory-mapped models with `memmap2` diff --git a/crates/ruvector-tiny-dancer-core/src/feature_engineering.rs b/crates/ruvector-tiny-dancer-core/src/feature_engineering.rs index 13cec7a1a1..67524542c4 100644 --- a/crates/ruvector-tiny-dancer-core/src/feature_engineering.rs +++ b/crates/ruvector-tiny-dancer-core/src/feature_engineering.rs @@ -5,6 +5,7 @@ use crate::error::{Result, TinyDancerError}; use crate::types::Candidate; use chrono::Utc; +#[cfg(all(not(feature = "lattice-simd"), feature = "simd-simsimd"))] use simsimd::SpatialSimilarity; /// Feature vector for a candidate @@ -130,7 +131,13 @@ impl FeatureEngineer { .collect() } - /// Compute cosine similarity using SIMD-optimized simsimd + /// Cosine similarity between two equal-length vectors. + /// + /// The kernel is feature-selected and exactly one arm compiles. All three + /// agree on the boundary conventions, which `backend_matches_scalar_reference` + /// checks: two all-zero vectors score `1.0`, and a zero vector against a + /// non-zero one scores `0.0`. Those are the values this path has always + /// returned, and a backend swap is not the place to change them. fn cosine_similarity(&self, a: &[f32], b: &[f32]) -> Result { if a.len() != b.len() { return Err(TinyDancerError::InvalidInput(format!( @@ -140,12 +147,32 @@ impl FeatureEngineer { ))); } - // Use simsimd for SIMD-accelerated cosine similarity - let similarity = f32::cosine(a, b) - .ok_or_else(|| TinyDancerError::FeatureError("Cosine similarity failed".to_string()))?; + #[cfg(feature = "lattice-simd")] + { + // Returns similarity directly, so there is no `1 - x` here. It also + // reports 0.0 for a zero norm, which collides with the genuine + // "orthogonal" answer, so the all-zero case is separated out on the + // cold path to keep the conventions above. + let similarity = lattice_embed::simd::cosine_similarity(a, b); + if similarity == 0.0 && is_all_zero(a) && is_all_zero(b) { + return Ok(1.0); + } + Ok(similarity) + } + + #[cfg(all(not(feature = "lattice-simd"), feature = "simd-simsimd"))] + { + // `SpatialSimilarity::cosine` returns a DISTANCE, hence the `1 - x`. + let distance = f32::cosine(a, b).ok_or_else(|| { + TinyDancerError::FeatureError("Cosine similarity failed".to_string()) + })?; + Ok(1.0_f32 - distance as f32) + } - // Convert distance to similarity (simsimd returns distance as f64) - Ok(1.0_f32 - similarity as f32) + #[cfg(all(not(feature = "lattice-simd"), not(feature = "simd-simsimd")))] + { + Ok(scalar_cosine_similarity(a, b)) + } } /// Calculate recency score using exponential decay @@ -201,6 +228,36 @@ impl Default for FeatureEngineer { } } +/// Scalar cosine similarity. Serves as the `--no-default-features` backend and +/// as the reference `backend_matches_scalar_reference` checks the compiled +/// backend against, so it stays compiled even when a SIMD arm is selected. +/// +/// A zero norm is not an error here: two all-zero vectors score `1.0` and a +/// zero vector against a non-zero one scores `0.0`, matching the other arms. +/// The `min` mirrors the clip a distance-returning kernel applies at 0. +fn scalar_cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let mut dot = 0.0f32; + let mut norm_a = 0.0f32; + let mut norm_b = 0.0f32; + for (x, y) in a.iter().zip(b.iter()) { + dot += x * y; + norm_a += x * x; + norm_b += y * y; + } + + if norm_a == 0.0 && norm_b == 0.0 { + return 1.0; + } + if dot == 0.0 { + return 0.0; + } + (dot / (norm_a.sqrt() * norm_b.sqrt())).min(1.0) +} + +fn is_all_zero(v: &[f32]) -> bool { + v.iter().all(|x| *x == 0.0) +} + #[cfg(test)] mod tests { use super::*; @@ -233,6 +290,89 @@ mod tests { assert!((similarity - 1.0).abs() < 0.01); } + /// Deterministic vector source. A fixed LCG keeps the grid reproducible + /// across arms without pulling a seeded-RNG dependency into the comparison. + fn lcg(state: &mut u64) -> f32 { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((*state >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0 + } + + /// Single-pass f64 reference. Deliberately NOT `scalar_cosine_similarity`: + /// that one accumulates in f32 and is a backend in its own right, so using + /// it as the reference would let a shared reduction-order bug pass. + fn naive_cosine(a: &[f32], b: &[f32]) -> f32 { + let mut dot = 0.0f64; + let mut norm_a = 0.0f64; + let mut norm_b = 0.0f64; + for (x, y) in a.iter().zip(b.iter()) { + dot += f64::from(*x) * f64::from(*y); + norm_a += f64::from(*x) * f64::from(*x); + norm_b += f64::from(*y) * f64::from(*y); + } + (dot / (norm_a.sqrt() * norm_b.sqrt())) as f32 + } + + /// Checks whichever backend compiled in against the reference, so this + /// covers the SimSIMD path shipping today, the scalar path, and the lattice + /// path alike. Dimensions straddle 4/8/16-lane widths and their remainders. + #[test] + fn backend_matches_scalar_reference() { + let engineer = FeatureEngineer::new(); + const DIMS: [usize; 21] = [ + 1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 63, 64, 128, 384, 768, 1000, 1023, 1024, + ]; + + for &dim in DIMS.iter() { + for seed_base in 0..3u64 { + let mut state = seed_base.wrapping_mul(7919).wrapping_add(12345); + let a: Vec = (0..dim).map(|_| lcg(&mut state)).collect(); + let b: Vec = (0..dim).map(|_| lcg(&mut state)).collect(); + + let got = engineer.cosine_similarity(&a, &b).unwrap(); + let want = naive_cosine(&a, &b); + assert!( + (got - want).abs() <= 1e-4 * want.abs().max(1.0), + "cosine mismatch dim={dim} seed={seed_base}: got {got}, want {want}" + ); + } + } + } + + /// The boundary values every backend must agree on. These are conventions + /// rather than arithmetic, so they are asserted exactly. + #[test] + fn backend_boundary_conventions_agree() { + let engineer = FeatureEngineer::new(); + let zero = vec![0.0f32; 8]; + let ones = vec![1.0f32; 8]; + + assert_eq!( + engineer.cosine_similarity(&zero, &zero).unwrap(), + 1.0, + "two all-zero vectors" + ); + assert_eq!( + engineer.cosine_similarity(&zero, &ones).unwrap(), + 0.0, + "zero against non-zero" + ); + + let e1 = vec![1.0, 0.0, 0.0, 0.0]; + let e2 = vec![0.0, 1.0, 0.0, 0.0]; + assert_eq!( + engineer.cosine_similarity(&e1, &e2).unwrap(), + 0.0, + "orthogonal" + ); + + assert!( + engineer.cosine_similarity(&e1, &ones).is_err(), + "length mismatch must stay an error on every backend" + ); + } + #[test] fn test_recency_score() { let engineer = FeatureEngineer::new();