From 9c898f0f894403b1045858a8b6c0779de8efbd45 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Sun, 2 Aug 2026 09:53:28 -0400 Subject: [PATCH 1/4] perf(maxsim): compute each query token's norm once per document maxsim calls cosine for every (query token, document token) pair, and cosine recomputed the query token's norm inside each of those calls. For a document of m tokens the same norm was accumulated m times. Hoists it to once per query token. The inner loop now does two multiply-adds per dimension instead of three, plus one norm per query token amortised over the whole document. The result is bit-identical, not merely close: the norm is accumulated in the same order and multiplied into the same denominator either way. A new test asserts that bitwise against the pre-hoist formulation, written out in full rather than delegating to cosine so a bug in the shared helper cannot cancel on both sides. Two more tests pin the degenerate cases, a zero-magnitude query token and a zero-magnitude document token. No public API change: cosine keeps its signature and now delegates to the same helper maxsim uses. --- crates/ruvector-maxsim/src/score.rs | 122 +++++++++++++++++++++++++++- 1 file changed, 118 insertions(+), 4 deletions(-) diff --git a/crates/ruvector-maxsim/src/score.rs b/crates/ruvector-maxsim/src/score.rs index 1afbf4c816..50221f3302 100644 --- a/crates/ruvector-maxsim/src/score.rs +++ b/crates/ruvector-maxsim/src/score.rs @@ -15,16 +15,35 @@ use crate::types::Embedding; /// Returns 0.0 when either vector is zero-magnitude. #[inline] pub fn cosine(a: &[f32], b: &[f32]) -> f32 { + debug_assert_eq!(a.len(), b.len(), "dimension mismatch in cosine"); + cosine_with_lhs_norm(a, b, norm(a)) +} + +/// L2 norm of a vector. +#[inline] +fn norm(v: &[f32]) -> f32 { + let mut acc = 0.0_f32; + for &x in v.iter() { + acc += x * x; + } + acc.sqrt() +} + +/// Cosine similarity with the left vector's norm supplied by the caller. +/// +/// Accumulates in the same order as [`cosine`] and combines the same two +/// factors into the same denominator, so a caller that passes `norm(a)` gets +/// bit-identical results to calling [`cosine`] directly. +#[inline] +fn cosine_with_lhs_norm(a: &[f32], b: &[f32], norm_a: f32) -> f32 { debug_assert_eq!(a.len(), b.len(), "dimension mismatch in cosine"); let mut dot = 0.0_f32; - let mut na = 0.0_f32; let mut nb = 0.0_f32; for (&ai, &bi) in a.iter().zip(b.iter()) { dot += ai * bi; - na += ai * ai; nb += bi * bi; } - let denom = na.sqrt() * nb.sqrt(); + let denom = norm_a * nb.sqrt(); if denom < f32::EPSILON { 0.0 } else { @@ -35,13 +54,18 @@ pub fn cosine(a: &[f32], b: &[f32]) -> f32 { /// MaxSim score between a multi-vector query and a multi-vector document. /// /// Time: O(|query_vecs| * |doc_vecs| * D). +/// +/// Each query token's norm is computed once and reused across every document +/// token, rather than being recomputed inside each pairwise cosine. The inner +/// loop therefore does two multiply-adds per dimension instead of three. pub fn maxsim(query_vecs: &[Embedding], doc_vecs: &[Embedding]) -> f32 { query_vecs .iter() .map(|q| { + let norm_q = norm(q); doc_vecs .iter() - .map(|d| cosine(q, d)) + .map(|d| cosine_with_lhs_norm(q, d, norm_q)) .fold(f32::NEG_INFINITY, f32::max) }) .sum() @@ -109,4 +133,94 @@ mod tests { // Each query token matches exactly one doc token → sum = 2.0 assert!((s - 2.0).abs() < 1e-5, "expected ~2.0, got {s}"); } + + /// The pre-hoist cosine, kept verbatim: one fused loop, three + /// accumulators, the query norm recomputed for every pair. + /// + /// Written out rather than delegating to [`cosine`] so that the oracle + /// shares no code with what it checks. Delegating would make a bug in the + /// shared helper cancel on both sides and the comparison pass anyway. + fn cosine_pre_hoist(a: &[f32], b: &[f32]) -> f32 { + let mut dot = 0.0_f32; + let mut na = 0.0_f32; + let mut nb = 0.0_f32; + for (&ai, &bi) in a.iter().zip(b.iter()) { + dot += ai * bi; + na += ai * ai; + nb += bi * bi; + } + let denom = na.sqrt() * nb.sqrt(); + if denom < f32::EPSILON { + 0.0 + } else { + dot / denom + } + } + + /// The pre-hoist formulation, kept verbatim as the oracle. + fn maxsim_recomputing_query_norm(q: &[Embedding], d: &[Embedding]) -> f32 { + q.iter() + .map(|qv| { + d.iter() + .map(|dv| cosine_pre_hoist(qv, dv)) + .fold(f32::NEG_INFINITY, f32::max) + }) + .sum() + } + + fn vecs(count: usize, dim: usize, seed: u32) -> Vec { + let mut s = seed.wrapping_mul(2_654_435_761).wrapping_add(1); + let mut next = || { + s ^= s << 13; + s ^= s >> 17; + s ^= s << 5; + (s as f32 / u32::MAX as f32) * 2.0 - 1.0 + }; + (0..count) + .map(|_| (0..dim).map(|_| next()).collect()) + .collect() + } + + /// Hoisting the query norm must be exact, not merely close. + /// + /// The norm is accumulated in the same order and multiplied into the same + /// denominator either way, so any difference at all would mean the + /// refactor changed the arithmetic. Compared bitwise for that reason. + #[test] + fn hoisting_query_norm_is_bit_exact() { + for dim in [1usize, 3, 8, 16, 33, 128, 384] { + for (nq, nd) in [(1usize, 1usize), (1, 7), (5, 1), (4, 9)] { + let q = vecs(nq, dim, dim as u32); + let d = vecs(nd, dim, dim as u32 + 17); + let got = maxsim(&q, &d); + let want = maxsim_recomputing_query_norm(&q, &d); + assert_eq!( + got.to_bits(), + want.to_bits(), + "dim={dim} nq={nq} nd={nd}: {got} vs {want}" + ); + } + } + } + + /// A zero-magnitude query token must still take the degenerate branch. + #[test] + fn zero_query_token_scores_zero() { + let q = vec![vec![0.0_f32; 8]]; + let d = vecs(4, 8, 3); + assert_eq!(maxsim(&q, &d), 0.0); + assert_eq!(maxsim(&q, &d), maxsim_recomputing_query_norm(&q, &d)); + } + + /// A zero-magnitude document token must not poison the max. + #[test] + fn zero_doc_token_does_not_poison_max() { + let q = vecs(2, 8, 5); + let mut d = vecs(3, 8, 11); + d.push(vec![0.0_f32; 8]); + assert_eq!( + maxsim(&q, &d).to_bits(), + maxsim_recomputing_query_norm(&q, &d).to_bits() + ); + } } From 8c0d91a84e620fa2dfa8a57d172b3e1125beafe3 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:50:59 -0400 Subject: [PATCH 2/4] perf(maxsim): restore cosine to a single fused pass cosine used to be one traversal of the operands accumulating dot, ||a||^2 and ||b||^2 together. It was changed to compute norm(a) as a separate pass and hand it to cosine_with_lhs_norm, so maxsim could hoist the query norm out of its inner loop over document tokens. That hoist helps maxsim, which reuses a fixed query token's norm across many document tokens: FlatMaxSim benchmarks improved by roughly 20%. But bucket.rs and hnsw.rs/graph.rs call the public cosine directly, once per vector pair, on their centroid and graph-traversal phases. For them the two-pass cosine only adds a redundant traversal of the left operand, and BucketMaxSim/HnswMaxSim benchmarks regressed accordingly. cosine is restored to its original single fused pass (three accumulators, one zip over both operands). cosine_with_lhs_norm, norm, and maxsim are unchanged, so the hoist and its benefit stay intact for maxsim's own inner loop while every direct cosine caller is back to its original cost. The test oracle that hand-duplicated the fused body is retired now that it would be identical to the restored cosine; the comparison instead checks that maxsim's hoisted composition (norm + cosine_with_lhs_norm) agrees with calling the public cosine per pair, which is the invariant that actually matters post-hoist. Co-Authored-By: claude-flow --- crates/ruvector-maxsim/src/score.rs | 50 ++++++++++++++--------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/crates/ruvector-maxsim/src/score.rs b/crates/ruvector-maxsim/src/score.rs index 50221f3302..3e1917b89a 100644 --- a/crates/ruvector-maxsim/src/score.rs +++ b/crates/ruvector-maxsim/src/score.rs @@ -16,7 +16,20 @@ use crate::types::Embedding; #[inline] pub fn cosine(a: &[f32], b: &[f32]) -> f32 { debug_assert_eq!(a.len(), b.len(), "dimension mismatch in cosine"); - cosine_with_lhs_norm(a, b, norm(a)) + let mut dot = 0.0_f32; + let mut na = 0.0_f32; + let mut nb = 0.0_f32; + for (&ai, &bi) in a.iter().zip(b.iter()) { + dot += ai * bi; + na += ai * ai; + nb += bi * bi; + } + let denom = na.sqrt() * nb.sqrt(); + if denom < f32::EPSILON { + 0.0 + } else { + dot / denom + } } /// L2 norm of a vector. @@ -134,35 +147,22 @@ mod tests { assert!((s - 2.0).abs() < 1e-5, "expected ~2.0, got {s}"); } - /// The pre-hoist cosine, kept verbatim: one fused loop, three - /// accumulators, the query norm recomputed for every pair. + /// MaxSim recomputed with the naive, un-hoisted formulation: every + /// pairwise score goes through the public [`cosine`], recomputing the + /// query token's norm on each call instead of once per query token. /// - /// Written out rather than delegating to [`cosine`] so that the oracle - /// shares no code with what it checks. Delegating would make a bug in the - /// shared helper cancel on both sides and the comparison pass anyway. - fn cosine_pre_hoist(a: &[f32], b: &[f32]) -> f32 { - let mut dot = 0.0_f32; - let mut na = 0.0_f32; - let mut nb = 0.0_f32; - for (&ai, &bi) in a.iter().zip(b.iter()) { - dot += ai * bi; - na += ai * ai; - nb += bi * bi; - } - let denom = na.sqrt() * nb.sqrt(); - if denom < f32::EPSILON { - 0.0 - } else { - dot / denom - } - } - - /// The pre-hoist formulation, kept verbatim as the oracle. + /// This is the property that matters after `cosine` went back to a + /// single fused pass: [`maxsim`]'s hoist (`norm(q)` once, then + /// [`cosine_with_lhs_norm`] per document token) must still agree with + /// calling the real `cosine` per pair. The two sides run genuinely + /// different code — one composes `norm` + `cosine_with_lhs_norm`, the + /// other calls fused `cosine` — so the comparison still guards a real + /// invariant instead of comparing a function with a hand-copy of itself. fn maxsim_recomputing_query_norm(q: &[Embedding], d: &[Embedding]) -> f32 { q.iter() .map(|qv| { d.iter() - .map(|dv| cosine_pre_hoist(qv, dv)) + .map(|dv| cosine(qv, dv)) .fold(f32::NEG_INFINITY, f32::max) }) .sum() From f706beba62044fa5935952f14f918f2859b4ab82 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:05:40 -0400 Subject: [PATCH 3/4] test(maxsim): count query-norm computations and pin zero-token scoring Add a thread-local test-only counter inside norm() (used exclusively by maxsim's hoisted query-norm computation) to assert exactly one query-norm computation per query token, independent of document count. Add direct zero-only-document and all-zero-document-tokens assertions (0.0, is_finite(), and equivalence to the un-hoisted oracle) so the denominator guard in cosine_with_lhs_norm is independently checkable rather than only implied by an equivalence test that a max-fold can mask. Co-Authored-By: Claude Sonnet 5 --- crates/ruvector-maxsim/src/score.rs | 60 +++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/crates/ruvector-maxsim/src/score.rs b/crates/ruvector-maxsim/src/score.rs index 3e1917b89a..2b6ae1de70 100644 --- a/crates/ruvector-maxsim/src/score.rs +++ b/crates/ruvector-maxsim/src/score.rs @@ -32,9 +32,21 @@ pub fn cosine(a: &[f32], b: &[f32]) -> f32 { } } +#[cfg(test)] +thread_local! { + /// Test-only seam counting calls to [`norm`], which `maxsim` uses + /// exclusively for the hoisted per-query-token norm. Thread-local (not a + /// shared atomic) so parallel test threads never pollute each other's + /// count; a given test's `maxsim` call runs synchronously on its own + /// thread, so this is race-free without any lock. + static QUERY_NORM_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + /// L2 norm of a vector. #[inline] fn norm(v: &[f32]) -> f32 { + #[cfg(test)] + QUERY_NORM_CALLS.with(|c| c.set(c.get() + 1)); let mut acc = 0.0_f32; for &x in v.iter() { acc += x * x; @@ -223,4 +235,52 @@ mod tests { maxsim_recomputing_query_norm(&q, &d).to_bits() ); } + + /// The hoist must retain its shape, not just its output: `maxsim` must + /// compute each query token's norm exactly once, no matter how many + /// document tokens it is scored against. More than one document token is + /// used deliberately — a formulation that recomputes the query norm once + /// per document token would tie with the hoisted one when `nd == 1`. + #[test] + fn hoist_computes_query_norm_once_per_token_regardless_of_doc_count() { + let q = vecs(3, 8, 51); + let d = vecs(5, 8, 59); + QUERY_NORM_CALLS.with(|c| c.set(0)); + let _ = maxsim(&q, &d); + let calls = QUERY_NORM_CALLS.with(|c| c.get()); + assert_eq!( + calls, + q.len(), + "expected exactly one query-norm computation per query token \ + regardless of document count, got {calls} for {} query tokens \ + against {} document tokens", + q.len(), + d.len() + ); + } + + /// A document containing ONLY a zero-magnitude token must score 0.0 and + /// finite, not NaN silently masked by the max-fold's NaN-ignoring + /// semantics. + #[test] + fn zero_only_document_scores_zero_and_finite() { + let q = vecs(2, 8, 41); + let d = vec![vec![0.0_f32; 8]]; + let got = maxsim(&q, &d); + assert_eq!(got, 0.0); + assert!(got.is_finite()); + assert_eq!(got, maxsim_recomputing_query_norm(&q, &d)); + } + + /// Every document token being zero-magnitude (not just one among + /// otherwise-normal tokens) must still resolve to 0.0 and finite. + #[test] + fn all_zero_document_tokens_score_zero_and_finite() { + let q = vecs(2, 8, 43); + let d = vec![vec![0.0_f32; 8]; 4]; + let got = maxsim(&q, &d); + assert_eq!(got, 0.0); + assert!(got.is_finite()); + assert_eq!(got, maxsim_recomputing_query_norm(&q, &d)); + } } From 0036f95c33c23474b6860614e0262b7b1770564d Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:37:40 -0400 Subject: [PATCH 4/4] fix(maxsim): keep base cosine semantics for mismatched dimensions maxsim's hoisted per-query-token norm is only bit-identical to the pre-hoist per-pair cosine when query and document tokens have equal length: norm(q) is taken over the full query, so combining it with a dot product zipped (and thus truncated) to a shorter document token silently changes the score. Fall back to a truncating fused cosine for mismatched-length pairs, matching the original single-pass cosine's truncating-zip behavior instead of the hoisted fast path's full-length norm. Co-Authored-By: Claude Sonnet 5 --- crates/ruvector-maxsim/src/score.rs | 74 +++++++++++++++++++++++++++-- 1 file changed, 70 insertions(+), 4 deletions(-) diff --git a/crates/ruvector-maxsim/src/score.rs b/crates/ruvector-maxsim/src/score.rs index 2b6ae1de70..7cd3e19554 100644 --- a/crates/ruvector-maxsim/src/score.rs +++ b/crates/ruvector-maxsim/src/score.rs @@ -58,7 +58,11 @@ fn norm(v: &[f32]) -> f32 { /// /// Accumulates in the same order as [`cosine`] and combines the same two /// factors into the same denominator, so a caller that passes `norm(a)` gets -/// bit-identical results to calling [`cosine`] directly. +/// bit-identical results to calling [`cosine`] directly — but only when `a` +/// and `b` have equal length: `norm_a` was computed over the full `a`, so a +/// truncated zip over `a` and `b` here would combine a full-length norm with +/// a truncated dot product. Callers with mismatched lengths must use +/// [`cosine`] (or [`cosine_truncating`]) instead. #[inline] fn cosine_with_lhs_norm(a: &[f32], b: &[f32], norm_a: f32) -> f32 { debug_assert_eq!(a.len(), b.len(), "dimension mismatch in cosine"); @@ -76,13 +80,42 @@ fn cosine_with_lhs_norm(a: &[f32], b: &[f32], norm_a: f32) -> f32 { } } +/// Fused cosine without [`cosine`]'s equal-length precondition: both norms +/// and the dot product accumulate inside the same `zip`, so mismatched +/// lengths truncate to the shorter vector rather than panicking in debug +/// builds. This is [`maxsim`]'s fallback for query/document token pairs of +/// differing dimension, where the hoisted `norm(q)` (taken over the full +/// query) would otherwise be combined with a dot product truncated to the +/// document's length. +#[inline] +fn cosine_truncating(a: &[f32], b: &[f32]) -> f32 { + let mut dot = 0.0_f32; + let mut na = 0.0_f32; + let mut nb = 0.0_f32; + for (&ai, &bi) in a.iter().zip(b.iter()) { + dot += ai * bi; + na += ai * ai; + nb += bi * bi; + } + let denom = na.sqrt() * nb.sqrt(); + if denom < f32::EPSILON { + 0.0 + } else { + dot / denom + } +} + /// MaxSim score between a multi-vector query and a multi-vector document. /// /// Time: O(|query_vecs| * |doc_vecs| * D). /// /// Each query token's norm is computed once and reused across every document -/// token, rather than being recomputed inside each pairwise cosine. The inner -/// loop therefore does two multiply-adds per dimension instead of three. +/// token of matching dimension, rather than being recomputed inside each +/// pairwise cosine. The inner loop therefore does two multiply-adds per +/// dimension instead of three. Pairs whose query and document tokens have +/// different lengths fall back to [`cosine_truncating`], since the hoisted +/// query norm is only bit-identical to per-pair cosine when both operands +/// are the same length. pub fn maxsim(query_vecs: &[Embedding], doc_vecs: &[Embedding]) -> f32 { query_vecs .iter() @@ -90,7 +123,13 @@ pub fn maxsim(query_vecs: &[Embedding], doc_vecs: &[Embedding]) -> f32 { let norm_q = norm(q); doc_vecs .iter() - .map(|d| cosine_with_lhs_norm(q, d, norm_q)) + .map(|d| { + if q.len() == d.len() { + cosine_with_lhs_norm(q, d, norm_q) + } else { + cosine_truncating(q, d) + } + }) .fold(f32::NEG_INFINITY, f32::max) }) .sum() @@ -283,4 +322,31 @@ mod tests { assert!(got.is_finite()); assert_eq!(got, maxsim_recomputing_query_norm(&q, &d)); } + + /// A query token longer than the document token it is scored against + /// must fall back to the pre-hoist, truncating-zip cosine rather than + /// combining a norm taken over the full query with a dot product + /// truncated to the document's length. q=[3,4,12] vs d=[3,4]: the + /// truncating base semantics give dot=25, na=25 (truncated to 2 terms), + /// nb=25, so 25/(5*5)=1.0 exactly. The hoisted fast path would instead + /// divide by norm(q)=13, giving 25/(13*5)=0.3846.... + #[test] + fn mismatched_query_longer_scores_via_truncating_fallback() { + let q = vec![vec![3.0_f32, 4.0, 12.0]]; + let d = vec![vec![3.0_f32, 4.0]]; + let got = maxsim(&q, &d); + assert_eq!(got, 1.0, "expected exact 1.0, got {got}"); + } + + /// The opposite mismatch direction — document token longer than the + /// query token — must also route through the truncating fallback and + /// agree with calling it directly on the same pair. + #[test] + fn mismatched_document_longer_matches_truncating_cosine() { + let q = vec![vec![3.0_f32, 4.0]]; + let d = vec![vec![3.0_f32, 4.0, 12.0]]; + let got = maxsim(&q, &d); + let want = cosine_truncating(&q[0], &d[0]); + assert_eq!(got, want); + } }