From 9f9d29e62a22bece29ee7c459e2c1192ac99a6a3 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:27:11 -0400 Subject: [PATCH 1/2] feat(tiny-dancer-core): make the cosine backend selectable The candidate-scoring cosine kernel was wired directly to one library through a hard dependency, and the crate had no `[features]` section at all, so there was no way to build it against anything else or without a SIMD library. This adds a feature axis with three mutually exclusive arms and leaves the default build on the existing backend: default (simd-simsimd) -> SimSIMD, unchanged lattice-simd -> lattice-embed kernels --no-default-features -> scalar The two SIMD libraries disagree on what cosine returns: one gives a distance and the other a similarity, so the adapters differ by the `1 - x` and the arms are not interchangeable text. They also disagree on zero norms. The conventions this path has always returned are preserved on every arm: two all-zero vectors score 1.0, a zero vector against a non-zero one scores 0.0, and a length mismatch stays an error. `backend_matches_scalar_reference` checks whichever backend compiled in against an f64 single-pass reference across 21 dimensions straddling 4/8/16-lane widths and their remainders, so it covers the backend shipping today as well as the new ones. Its reference is deliberately not the crate's own scalar function, which is itself an f32 4-accumulator backend: using it would let a shared reduction-order bug pass. `backend_boundary_conventions_agree` pins the zero-norm and mismatch values exactly. lattice-embed requires Rust >= 1.93 and 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. --- Cargo.lock | 21 ++- crates/ruvector-tiny-dancer-core/Cargo.toml | 18 ++- .../src/feature_engineering.rs | 152 +++++++++++++++++- 3 files changed, 183 insertions(+), 8 deletions(-) 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/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(); From dacd46f9bdaae5f3cf620c80d4fd55fca9a3767e Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:41:54 -0400 Subject: [PATCH 2/2] docs(tiny-dancer): document the selectable distance backends --- crates/ruvector-tiny-dancer-core/README.md | 26 ++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) 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`