From 00d14cfa93fe745e31d565e1e6a2a764e7466030 Mon Sep 17 00:00:00 2001 From: David Snelling Date: Wed, 27 May 2026 12:43:29 -0700 Subject: [PATCH] feat: hook native SQ8 distance provider into HNSW reranking Makes distanceSQ8 swappable via setSQ8DistanceImplementation and wires brainy.ts to install a registered 'distance:sq8' provider (e.g. cortex's Rust SIMD), with the JS distanceSQ8Js as the default fallback. The provider signature is byte-compatible with the native cosineDistanceSq8, so HNSW SQ8 reranking transparently uses native distance when cortex is loaded. Native==JS numeric equivalence is asserted by the cross-language parity suite. --- src/brainy.ts | 11 +++++ src/utils/vectorQuantization.ts | 60 +++++++++++++++++++++++++++- tests/unit/hnsw/quantization.test.ts | 45 ++++++++++++++++++++- 3 files changed, 114 insertions(+), 2 deletions(-) diff --git a/src/brainy.ts b/src/brainy.ts index 24c91fbf..368f6944 100644 --- a/src/brainy.ts +++ b/src/brainy.ts @@ -499,6 +499,17 @@ export class Brainy implements BrainyInterface { setMsgpackImplementation(msgpackProvider) } + // Provider: native SQ8 approximate distance (e.g. cortex's Rust SIMD) — swaps + // the JS quantized-distance used in HNSW SQ8 reranking. Signature is + // byte-compatible with the JS distanceSQ8; falls back to JS when absent. + const sq8DistanceProvider = this.pluginRegistry.getProvider< + (a: Uint8Array, aMin: number, aMax: number, b: Uint8Array, bMin: number, bMax: number) => number + >('distance:sq8') + if (sq8DistanceProvider) { + const { setSQ8DistanceImplementation } = await import('./utils/vectorQuantization.js') + setSQ8DistanceImplementation(sq8DistanceProvider) + } + // Provider: distance function (resolve BEFORE setupIndex — index uses this.distance) const nativeDistance = this.pluginRegistry.getProvider('distance') if (nativeDistance) { diff --git a/src/utils/vectorQuantization.ts b/src/utils/vectorQuantization.ts index 648e422a..36234152 100644 --- a/src/utils/vectorQuantization.ts +++ b/src/utils/vectorQuantization.ts @@ -117,7 +117,12 @@ export function dequantizeSQ8(quantized: Uint8Array, min: number, max: number): * @param bMax - Second vector codebook maximum * @returns Approximate cosine distance [0, 2] */ -export function distanceSQ8( +/** + * Reference JS implementation of {@link distanceSQ8}. Exported so callers can + * explicitly restore the default after swapping in a native implementation, and + * so cross-language parity tests can compare native output against this baseline. + */ +export function distanceSQ8Js( a: Uint8Array, aMin: number, aMax: number, @@ -185,6 +190,59 @@ export function distanceSQ8( return 1 - Math.max(-1, Math.min(1, cosine)) } +/** Signature of the SQ8 approximate-distance function (quantized cosine distance). */ +export type SQ8DistanceFn = ( + a: Uint8Array, + aMin: number, + aMax: number, + b: Uint8Array, + bMin: number, + bMax: number +) => number + +/** + * Active SQ8 distance implementation. Defaults to the JS `distanceSQ8Js`; swapped to + * a native SIMD implementation by {@link setSQ8DistanceImplementation} when a cortex + * `distance:sq8` provider is registered. The native implementation is byte-compatible + * (same quantized cosine formula), so results match within float tolerance. + */ +let activeSQ8Distance: SQ8DistanceFn = distanceSQ8Js + +/** + * Approximate cosine distance between two SQ8-quantized vectors, dispatched to the + * active implementation (JS by default, native when a `distance:sq8` provider is + * registered). Used in the HNSW SQ8 reranking hot path. + * + * @param a - First quantized vector. + * @param aMin - First vector's codebook minimum. + * @param aMax - First vector's codebook maximum. + * @param b - Second quantized vector. + * @param bMin - Second vector's codebook minimum. + * @param bMax - Second vector's codebook maximum. + * @returns Approximate cosine distance in [0, 2]. + */ +export function distanceSQ8( + a: Uint8Array, + aMin: number, + aMax: number, + b: Uint8Array, + bMin: number, + bMax: number +): number { + return activeSQ8Distance(a, aMin, aMax, b, bMin, bMax) +} + +/** + * Replace the SQ8 approximate-distance implementation at runtime (e.g. cortex's Rust + * SIMD `distance:sq8`). Called by `brainy.ts` when the provider is registered. Pass + * {@link distanceSQ8Js} to restore the JS default. + * + * @param fn - Native implementation; must match {@link SQ8DistanceFn} byte-for-byte. + */ +export function setSQ8DistanceImplementation(fn: SQ8DistanceFn): void { + activeSQ8Distance = fn +} + /** * Serialize an SQ8 quantized vector to a compact binary format. * diff --git a/tests/unit/hnsw/quantization.test.ts b/tests/unit/hnsw/quantization.test.ts index f6a59921..8ea50d63 100644 --- a/tests/unit/hnsw/quantization.test.ts +++ b/tests/unit/hnsw/quantization.test.ts @@ -10,12 +10,14 @@ * - Config defaults: quantization disabled by default (no behavior change) */ -import { describe, it, expect, beforeEach } from 'vitest' +import { describe, it, expect, beforeEach, afterEach } from 'vitest' import { v4 as uuidv4 } from 'uuid' import { quantizeSQ8, dequantizeSQ8, distanceSQ8, + distanceSQ8Js, + setSQ8DistanceImplementation, serializeSQ8, deserializeSQ8 } from '../../../src/utils/vectorQuantization.js' @@ -28,6 +30,47 @@ function randomVector(dim: number): number[] { return Array.from({ length: dim }, () => Math.random() * 2 - 1) } +describe('SQ8 distance provider hook (setSQ8DistanceImplementation)', () => { + const q1 = quantizeSQ8([0.1, 0.5, -0.3, 0.8]) + const q2 = quantizeSQ8([0.2, 0.4, -0.1, 0.7]) + + // Always restore the JS default so a swapped impl never leaks into other tests. + afterEach(() => setSQ8DistanceImplementation(distanceSQ8Js)) + + it('dispatches to the JS implementation by default (identical result)', () => { + const viaDispatcher = distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max) + const viaJs = distanceSQ8Js(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max) + expect(viaDispatcher).toBe(viaJs) + expect(viaDispatcher).toBeGreaterThanOrEqual(0) + expect(viaDispatcher).toBeLessThanOrEqual(2) + }) + + it('routes through a swapped implementation with the cortex 6-arg signature', () => { + const calls: Array<[Uint8Array, number, number, Uint8Array, number, number]> = [] + setSQ8DistanceImplementation((a, aMin, aMax, b, bMin, bMax) => { + calls.push([a, aMin, aMax, b, bMin, bMax]) + return 0.4242 + }) + const result = distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max) + expect(result).toBe(0.4242) + expect(calls).toHaveLength(1) + // Exactly the arguments cortex's native cosineDistanceSq8(a, aMin, aMax, b, bMin, bMax) expects. + expect(calls[0][0]).toBe(q1.quantized) + expect(calls[0][1]).toBe(q1.min) + expect(calls[0][3]).toBe(q2.quantized) + expect(calls[0][5]).toBe(q2.max) + }) + + it('restores the JS default when set back', () => { + setSQ8DistanceImplementation(() => -999) + expect(distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)).toBe(-999) + setSQ8DistanceImplementation(distanceSQ8Js) + expect(distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)).toBe( + distanceSQ8Js(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max) + ) + }) +}) + // Helper: cosine distance for float32 vectors (reference implementation) function cosineDistance(a: number[], b: number[]): number { let dot = 0