Merge feat/int8-kernel: int8 index is faster than f32 on AVX2
CI / test (push) Failing after 3s

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:
osobh
2026-09-20 17:32:40 -07:00
co-authored by Claude Opus 5
8 changed files with 188 additions and 33 deletions
+25 -5
View File
@@ -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 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 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 is done automatically whenever the index is quantised.
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 **On AVX2 this costs nothing — it pays.** The first measurement of this put
recall. It is **off by default**: the right side of that trade depends on the cost at ~13% of QPS and ~16% of build time, but that compared a scalar
whether the deployment is short of memory or short of CPU. 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 A measurement trap worth recording: the synthetic `clustered` generator in the
`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so `clawhdf5-ann` tests draws clusters far tighter than any real embedding, so
+14
View File
@@ -2,6 +2,20 @@
## Unreleased ## 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 ### Correctness
- `clawhdf5-format`: **datasets indexed by an Extensible Array returned wrong - `clawhdf5-format`: **datasets indexed by an Extensible Array returned wrong
data beyond their first few dozen chunks.** One unlimited dimension gives a data beyond their first few dozen chunks.** One unlimited dimension gives a
+4 -2
View File
@@ -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 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 at 100K); because quantised distances are approximate and `ef` cannot
compensate, the query path then re-scores the candidate pool against the 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 exact embeddings, which holds recall at the f32 index's level. On AVX2 it is
~13% of QPS. `hybrid_search` keeps one incremental BM25 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 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 activation boosts are persisted by the next checkpoint (or on drop), not per
query. Measure any search-path change with query. Measure any search-path change with
+5 -2
View File
@@ -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 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 (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 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 embeddings the store already holds, which keeps recall at the `f32` index's
13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy". 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 | | `parallel` | no | Rayon parallel search |
| `fast-math` | no | BLAS matrix-vector multiply | | `fast-math` | no | BLAS matrix-vector multiply |
| `accelerate` | no | Apple Accelerate / AMX (macOS) | | `accelerate` | no | Apple Accelerate / AMX (macOS) |
+49
View File
@@ -25,6 +25,55 @@ unsafe fn hsum_256(v: __m256) -> f32 {
_mm_cvtss_f32(result) _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. /// AVX2 dot product for f32 slices.
/// ///
/// # Safety /// # Safety
+56
View File
@@ -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. /// Compute the L2 norm (magnitude) of a vector.
pub fn vector_norm(v: &[f32]) -> f32 { pub fn vector_norm(v: &[f32]) -> f32 {
dot_product(v, v).sqrt() 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);
}
}
+30
View File
@@ -140,3 +140,33 @@ fn f16_to_f32_soft(h: u16) -> f32 {
f32::from_bits(f32_bits) 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
}
+5 -24
View File
@@ -358,31 +358,12 @@ enum Query {
Int8(Vec<i8>, f32), Int8(Vec<i8>, f32),
} }
/// Sum of products, widened so it cannot overflow: `dim` terms of at most /// Sum of products, widened so it cannot overflow. Runtime-dispatched to the
/// `127 * 127`, so `i32` suffices for any realistic dimension. /// 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 { fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
// Four independent accumulators over 32-lane blocks: the widening product clawhdf5_accel::dot_i8(a, b)
// 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
} }
/// Magic for [`HnswIndex::graph_to_bytes`]. /// Magic for [`HnswIndex::graph_to_bytes`].