The quantised HNSW index was measured against f32 with a scalar loop on one side and clawhdf5-accel's AVX2 kernels on the other, so the ~13% throughput cost attributed to quantisation was a missing kernel. With clawhdf5_accel::dot_i8 in place, at equal recall the quantised index answers 1.63x the queries per second and builds 1.8x faster, holding a quarter of the vectors. Still off by default, now because the kernel is AVX2-only and aarch64 falls back to scalar. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
+25
-5
@@ -88,11 +88,31 @@ back**, because the loss is in the distances rather than in the graph
|
||||
|
||||
Re-scoring closes the gap: the store already holds the exact embeddings, so
|
||||
the query path re-scores the candidate pool against them before fusion. That
|
||||
is done automatically whenever the index is quantised. What it costs is
|
||||
throughput — about 13% of QPS and 16% of build time at 100 000 x 384. So the
|
||||
setting trades ~13% of query speed for ~36% of the process's memory at equal
|
||||
recall. It is **off by default**: the right side of that trade depends on
|
||||
whether the deployment is short of memory or short of CPU.
|
||||
is done automatically whenever the index is quantised.
|
||||
|
||||
**On AVX2 this costs nothing — it pays.** The first measurement of this put
|
||||
the cost at ~13% of QPS and ~16% of build time, but that compared a scalar
|
||||
int8 loop against `clawhdf5-accel`'s hand-written AVX2 kernels for `f32`:
|
||||
the gap was a missing kernel, not a property of int8. With
|
||||
`clawhdf5_accel::dot_i8` (AVX2: sign-extend to `i16`, then `madd_epi16`),
|
||||
medians of three alternating runs at N = 100 000, same binary:
|
||||
|
||||
| | f32 | int8 | int8 + re-score |
|
||||
|---|---:|---:|---:|
|
||||
| build | 3197 ms | **1778 ms** | 1826 ms |
|
||||
| QPS at ef = 64 | 13 399 | 29 195 | **21 848** |
|
||||
| recall@10 at ef = 64 | 0.9945 | 0.9625 | **0.9940** |
|
||||
|
||||
So at equal recall the quantised index answers **1.63x as many queries per
|
||||
second**, builds **1.8x faster**, and holds a quarter of the vectors. (Compare
|
||||
only at equal `ef`: with re-scoring the harness raises `ef` to at least the
|
||||
candidate pool, so the `ef = 16` and `ef = 32` rows are not like-for-like.)
|
||||
|
||||
It is still **off by default**, for portability rather than performance: the
|
||||
int8 kernel is AVX2-only, and on aarch64 — including `clawhdf5-android` — it
|
||||
falls back to the scalar loop, where the original trade still applies. A NEON
|
||||
kernel would remove that caveat. On an x86-64 deployment, turning it on is a
|
||||
win on every axis measured.
|
||||
|
||||
A measurement trap worth recording: the synthetic `clustered` generator in the
|
||||
`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so
|
||||
|
||||
@@ -2,6 +2,20 @@
|
||||
|
||||
## Unreleased
|
||||
|
||||
### Performance
|
||||
- `clawhdf5-accel`: **`dot_i8`, a runtime-dispatched int8 dot product** (AVX2:
|
||||
sign-extend each half to `i16`, then `madd_epi16`; scalar fallback
|
||||
elsewhere). The quantised HNSW index used a scalar loop while the `f32` path
|
||||
it was measured against ran AVX2, so the ~13% throughput cost recorded for
|
||||
`MemoryConfig::quantized_index` was a missing kernel rather than a property
|
||||
of int8. With the kernel, at N = 100 000 x 384 and equal recall, the
|
||||
quantised index answers **1.63x as many queries per second** (21 848 vs
|
||||
13 399 at ef=64, recall 0.9940 vs 0.9945) and builds **1.8x faster** (1778
|
||||
vs 3197 ms) — on top of holding a quarter of the vectors. Medians of three
|
||||
alternating runs. It remains off by default only because the kernel is
|
||||
AVX2-only and aarch64 falls back to the scalar loop. Integer arithmetic, so
|
||||
the SIMD path is tested to agree with scalar bit for bit.
|
||||
|
||||
### Correctness
|
||||
- `clawhdf5-format`: **datasets indexed by an Extensible Array returned wrong
|
||||
data beyond their first few dozen chunks.** One unlimited dimension gives a
|
||||
|
||||
@@ -44,8 +44,10 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
|
||||
which roughly halves a loaded store's memory (2.72x -> 1.74x the raw vectors
|
||||
at 100K); because quantised distances are approximate and `ef` cannot
|
||||
compensate, the query path then re-scores the candidate pool against the
|
||||
exact embeddings, which holds recall at the f32 index's level and costs
|
||||
~13% of QPS. `hybrid_search` keeps one incremental BM25
|
||||
exact embeddings, which holds recall at the f32 index's level. On AVX2 it is
|
||||
also 1.63x the QPS and 1.8x the build speed (`clawhdf5_accel::dot_i8`); it
|
||||
stays off by default only because that kernel is AVX2-only and aarch64 falls
|
||||
back to scalar. `hybrid_search` keeps one incremental BM25
|
||||
index for the life of the store and never writes the store: Hebbian
|
||||
activation boosts are persisted by the next checkpoint (or on drop), not per
|
||||
query. Measure any search-path change with
|
||||
|
||||
@@ -437,8 +437,11 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
|
||||
copy of the embeddings as `i8`, roughly halving a loaded store's memory
|
||||
(2.72x -> 1.74x the raw vectors at 100k x 384). Quantised distances are
|
||||
approximate, so the query path re-scores the candidate pool against the exact
|
||||
embeddings the store already holds — recall matches the `f32` index, at about
|
||||
13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy".
|
||||
embeddings the store already holds, which keeps recall at the `f32` index's
|
||||
level. On AVX2 it is also **faster** — 1.63x the queries per second and 1.8x
|
||||
the build speed at equal recall — because the int8 kernel is SIMD too. It
|
||||
stays off by default only because that kernel is AVX2-only and aarch64 falls
|
||||
back to a scalar loop. See `BENCHMARKS.md`, "Quantising the index copy".
|
||||
| `parallel` | no | Rayon parallel search |
|
||||
| `fast-math` | no | BLAS matrix-vector multiply |
|
||||
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
|
||||
|
||||
@@ -25,6 +25,55 @@ unsafe fn hsum_256(v: __m256) -> f32 {
|
||||
_mm_cvtss_f32(result)
|
||||
}
|
||||
|
||||
/// AVX2 dot product of two `i8` slices, widened to `i32`.
|
||||
///
|
||||
/// Each 16-byte half is sign-extended to sixteen `i16` lanes and multiplied
|
||||
/// pairwise with `madd_epi16`, which sums adjacent products straight into
|
||||
/// eight `i32` lanes — the widening that an autovectorised scalar loop does
|
||||
/// in several shuffles is one instruction here. A pair sum is at most
|
||||
/// `2 * 127 * 127`, far inside `i32`.
|
||||
///
|
||||
/// # Safety
|
||||
/// Caller must verify is_x86_feature_detected!("avx2").
|
||||
// SAFETY: Caller must have verified AVX2 via is_x86_feature_detected!.
|
||||
#[target_feature(enable = "avx2")]
|
||||
pub unsafe fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||
// SAFETY: Caller guarantees AVX2 is available per the # Safety contract;
|
||||
// every load reads 32 bytes at an index checked against `len` first.
|
||||
unsafe {
|
||||
assert_eq!(a.len(), b.len());
|
||||
let len = a.len();
|
||||
let mut i = 0;
|
||||
let mut acc0 = _mm256_setzero_si256();
|
||||
let mut acc1 = _mm256_setzero_si256();
|
||||
|
||||
while i + 32 <= len {
|
||||
let va = _mm256_loadu_si256(a.as_ptr().add(i).cast());
|
||||
let vb = _mm256_loadu_si256(b.as_ptr().add(i).cast());
|
||||
let a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(va));
|
||||
let b_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(vb));
|
||||
let a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(va, 1));
|
||||
let b_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(vb, 1));
|
||||
acc0 = _mm256_add_epi32(acc0, _mm256_madd_epi16(a_lo, b_lo));
|
||||
acc1 = _mm256_add_epi32(acc1, _mm256_madd_epi16(a_hi, b_hi));
|
||||
i += 32;
|
||||
}
|
||||
|
||||
// Horizontal sum of the eight i32 lanes.
|
||||
let v = _mm256_add_epi32(acc0, acc1);
|
||||
let s128 = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1));
|
||||
let s64 = _mm_add_epi32(s128, _mm_unpackhi_epi64(s128, s128));
|
||||
let s32 = _mm_add_epi32(s64, _mm_shuffle_epi32(s64, 0b01));
|
||||
let mut sum = _mm_cvtsi128_si32(s32);
|
||||
|
||||
while i < len {
|
||||
sum += i32::from(a[i]) * i32::from(b[i]);
|
||||
i += 1;
|
||||
}
|
||||
sum
|
||||
}
|
||||
}
|
||||
|
||||
/// AVX2 dot product for f32 slices.
|
||||
///
|
||||
/// # Safety
|
||||
|
||||
@@ -122,6 +122,23 @@ pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
|
||||
}
|
||||
}
|
||||
|
||||
/// Dot product of two `i8` slices, widened to `i32`.
|
||||
///
|
||||
/// The kernel behind int8-quantised vector search. Uses the AVX2 path
|
||||
/// whenever AVX2 is present — including on AVX-512 machines, where it is
|
||||
/// what the f32 kernels use too on a default build.
|
||||
pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||
match detect_backend() {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
// SAFETY: both variants imply AVX2 was detected at runtime (the
|
||||
// AVX-512 backend is only selected on CPUs that also have AVX2).
|
||||
Backend::Avx2 | Backend::Avx512 if is_x86_feature_detected!("avx2") => unsafe {
|
||||
avx2::dot_i8(a, b)
|
||||
},
|
||||
_ => scalar::dot_i8(a, b),
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the L2 norm (magnitude) of a vector.
|
||||
pub fn vector_norm(v: &[f32]) -> f32 {
|
||||
dot_product(v, v).sqrt()
|
||||
@@ -713,3 +730,42 @@ mod tests {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod dot_i8_tests {
|
||||
use super::*;
|
||||
|
||||
fn codes(n: usize, seed: u64) -> Vec<i8> {
|
||||
let mut state = seed;
|
||||
(0..n)
|
||||
.map(|_| {
|
||||
state = state
|
||||
.wrapping_mul(6_364_136_223_846_793_005)
|
||||
.wrapping_add(1_442_695_040_888_963_407);
|
||||
// Full range, including the extremes.
|
||||
((state >> 56) as u8) as i8
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatched_kernel_matches_scalar_exactly() {
|
||||
// Integer arithmetic: the SIMD path must agree bit for bit, at every
|
||||
// length — including ones that are not multiples of the 32-byte block,
|
||||
// which exercise the tail.
|
||||
for len in [0, 1, 7, 31, 32, 33, 63, 64, 100, 384, 385, 1536] {
|
||||
let a = codes(len, 1 + len as u64);
|
||||
let b = codes(len, 1000 + len as u64);
|
||||
assert_eq!(dot_i8(&a, &b), scalar::dot_i8(&a, &b), "len {len}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extremes_do_not_overflow() {
|
||||
// -128 * -128 is the largest product; a long run of it must still fit.
|
||||
let a = vec![-128i8; 4096];
|
||||
assert_eq!(dot_i8(&a, &a), 4096 * 128 * 128);
|
||||
let b = vec![127i8; 4096];
|
||||
assert_eq!(dot_i8(&a, &b), -4096 * 128 * 127);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,3 +140,33 @@ fn f16_to_f32_soft(h: u16) -> f32 {
|
||||
|
||||
f32::from_bits(f32_bits)
|
||||
}
|
||||
|
||||
/// Dot product of two `i8` slices, widened to `i32`.
|
||||
///
|
||||
/// `dim` terms of at most `127 * 127` fit an `i32` for any realistic
|
||||
/// dimension (over 130 000 terms before overflow is possible).
|
||||
pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||
assert_eq!(a.len(), b.len());
|
||||
// Four independent accumulators over 32-lane blocks: the widening product
|
||||
// has to sit in a fixed-length chunk for the vectoriser to see it, and the
|
||||
// separate accumulators keep it off one dependency chain.
|
||||
const LANE: usize = 8;
|
||||
let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>();
|
||||
let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>();
|
||||
let mut acc = [0i32; 4];
|
||||
for (x, y) in a_blocks.iter().zip(b_blocks) {
|
||||
for (lane, slot) in acc.iter_mut().enumerate() {
|
||||
let mut sum = 0i32;
|
||||
for k in 0..LANE {
|
||||
sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]);
|
||||
}
|
||||
*slot += sum;
|
||||
}
|
||||
}
|
||||
let tail: i32 = a_tail
|
||||
.iter()
|
||||
.zip(b_tail)
|
||||
.map(|(&x, &y)| i32::from(x) * i32::from(y))
|
||||
.sum();
|
||||
acc[0] + acc[1] + acc[2] + acc[3] + tail
|
||||
}
|
||||
|
||||
@@ -358,31 +358,12 @@ enum Query {
|
||||
Int8(Vec<i8>, f32),
|
||||
}
|
||||
|
||||
/// Sum of products, widened so it cannot overflow: `dim` terms of at most
|
||||
/// `127 * 127`, so `i32` suffices for any realistic dimension.
|
||||
/// Sum of products, widened so it cannot overflow. Runtime-dispatched to the
|
||||
/// same SIMD backend as the f32 kernels, so the two storages are compared on
|
||||
/// equal terms.
|
||||
#[inline]
|
||||
fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
|
||||
// Four independent accumulators over 32-lane blocks: the widening product
|
||||
// has to sit in a fixed-length chunk for the vectoriser to see it, and the
|
||||
// separate accumulators keep it off one dependency chain.
|
||||
const LANE: usize = 8;
|
||||
let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>();
|
||||
let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>();
|
||||
let mut acc = [0i32; 4];
|
||||
for (x, y) in a_blocks.iter().zip(b_blocks) {
|
||||
for (lane, slot) in acc.iter_mut().enumerate() {
|
||||
let mut sum = 0i32;
|
||||
for k in 0..LANE {
|
||||
sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]);
|
||||
}
|
||||
*slot += sum;
|
||||
}
|
||||
}
|
||||
let tail: i32 = a_tail
|
||||
.iter()
|
||||
.zip(b_tail)
|
||||
.map(|(&x, &y)| i32::from(x) * i32::from(y))
|
||||
.sum();
|
||||
acc[0] + acc[1] + acc[2] + acc[3] + tail
|
||||
clawhdf5_accel::dot_i8(a, b)
|
||||
}
|
||||
|
||||
/// Magic for [`HnswIndex::graph_to_bytes`].
|
||||
|
||||
Reference in New Issue
Block a user