diff --git a/benches/nanobeir_baseline.json b/benches/nanobeir_baseline.json index 0a8799b..d9ebf8d 100644 --- a/benches/nanobeir_baseline.json +++ b/benches/nanobeir_baseline.json @@ -118,5 +118,20 @@ "NanoSCIDOCS": 0.259793, "NanoSciFact": 0.578403, "NanoTouche2020": 0.475323 + }, + "hotchpotch/bekko-static-embedding-v1": { + "NanoArguAna": 0.481588, + "NanoClimateFEVER": 0.346348, + "NanoDBPedia": 0.56677, + "NanoFEVER": 0.733361, + "NanoFiQA2018": 0.363626, + "NanoHotpotQA": 0.724093, + "NanoMSMARCO": 0.448949, + "NanoNFCorpus": 0.357624, + "NanoNQ": 0.421227, + "NanoQuoraRetrieval": 0.932645, + "NanoSCIDOCS": 0.357047, + "NanoSciFact": 0.711873, + "NanoTouche2020": 0.523074 } } diff --git a/scripts/bench_compare.py b/scripts/bench_compare.py index d29b48d..be454a9 100755 --- a/scripts/bench_compare.py +++ b/scripts/bench_compare.py @@ -48,6 +48,7 @@ "minishlab/potion-retrieval-32M", "sentence-transformers/static-retrieval-mrl-en-v1", "minishlab/potion-multilingual-128M", + "hotchpotch/bekko-static-embedding-v1", str(REPO / "tests" / "fixtures" / "test-model-vocab-quantized"), ] MIN_ROUNDS_FOR_CI = 5 diff --git a/src/model.rs b/src/model.rs index 9b2a5b4..473caf2 100644 --- a/src/model.rs +++ b/src/model.rs @@ -617,14 +617,20 @@ impl StaticModel { Embeddings::I8(table) => self.sum_rows(table, ids, sum), }; let denom = cnt.max(1) as f32; - for x in sum.iter_mut() { - *x /= denom; - } if model.normalize { - let norm = sum.iter().map(|&v| v * v).sum::().sqrt().max(1e-12); + let mut sq = 0.0_f32; + for x in sum.iter_mut() { + *x /= denom; + sq += *x * *x; + } + let norm = sq.sqrt().max(1e-12); for x in sum.iter_mut() { *x /= norm; } + } else { + for x in sum.iter_mut() { + *x /= denom; + } } } @@ -650,15 +656,8 @@ impl StaticModel { where f32: From, { - let mut sum = [0.0_f32; D]; - let cnt = self.for_each_row(table, ids, |row, scale| { - let row: &[T; D] = row.try_into().expect("embedding row has D columns"); - for (s, &v) in sum.iter_mut().zip(row) { - *s += f32::from(v) * scale; - } - }); - out.copy_from_slice(&sum); - cnt + let sum: &mut [f32; D] = out.try_into().expect("output has D columns"); + self.sum_rows_blocked(table, ids, sum) } /// Weighted row sum for any embedding dimension. @@ -666,17 +665,36 @@ impl StaticModel { where f32: From, { - self.for_each_row(table, ids, |row, scale| { - for (s, &v) in sum.iter_mut().zip(row) { - *s += f32::from(v) * scale; + self.sum_rows_blocked(table, ids, sum) + } + + /// Weighted row sum that adds `BLOCK` rows per pass over `sum`, in token + /// order, so the result matches adding one row at a time. + #[inline(always)] + fn sum_rows_blocked(&self, table: &[T], ids: &[u32], sum: &mut [f32]) -> usize + where + f32: From, + { + let mut pending: [(&[T], f32); BLOCK] = [(&[], 0.0); BLOCK]; + let mut n = 0; + let cnt = self.for_each_row(table, ids, |row, scale| { + pending[n] = (row, scale); + n += 1; + if n == BLOCK { + add_rows(sum, &pending); + n = 0; } - }) + }); + for &(row, scale) in &pending[..n] { + add_row(sum, row, scale); + } + cnt } /// Call `f` with the embedding row and scale (token weight times row /// scale) of every non-unk token; returns the number of rows visited. #[inline(always)] - fn for_each_row(&self, table: &[T], ids: &[u32], mut f: impl FnMut(&[T], f32)) -> usize { + fn for_each_row<'t, T>(&self, table: &'t [T], ids: &[u32], mut f: impl FnMut(&'t [T], f32)) -> usize { let model = &self.inner; let dim = model.dim; let mut cnt = 0; @@ -713,6 +731,59 @@ impl StaticModel { } } +/// `sum += row * scale`. +#[inline(always)] +fn add_row(sum: &mut [f32], row: &[T], scale: f32) +where + f32: From, +{ + for (s, &v) in sum.iter_mut().zip(row) { + *s += f32::from(v) * scale; + } +} + +const BLOCK: usize = 8; + +/// `sum += r0 * s0 + r1 * s1 + ...` over BLOCK rows, added left to right. +#[inline(always)] +fn add_rows(sum: &mut [f32], rows: &[(&[T], f32); BLOCK]) +where + f32: From, +{ + let n = sum.len(); + let [ + (r0, s0), + (r1, s1), + (r2, s2), + (r3, s3), + (r4, s4), + (r5, s5), + (r6, s6), + (r7, s7), + ] = *rows; + let (r0, r1, r2, r3, r4, r5, r6, r7) = ( + &r0[..n], + &r1[..n], + &r2[..n], + &r3[..n], + &r4[..n], + &r5[..n], + &r6[..n], + &r7[..n], + ); + for j in 0..n { + sum[j] = sum[j] + + f32::from(r0[j]) * s0 + + f32::from(r1[j]) * s1 + + f32::from(r2[j]) * s2 + + f32::from(r3[j]) * s3 + + f32::from(r4[j]) * s4 + + f32::from(r5[j]) * s5 + + f32::from(r6[j]) * s6 + + f32::from(r7[j]) * s7; + } +} + /// Symmetric per-row int8 quantization: each row is divided by /// `max|x| / 127`, rounded, and clamped to [-127, 127]. fn quantize_rows(table: &[f32], dim: usize) -> (Vec, Vec) {