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.
This commit is contained in:
parent
e23361c488
commit
00d14cfa93
3 changed files with 114 additions and 2 deletions
|
|
@ -499,6 +499,17 @@ export class Brainy<T = any> implements BrainyInterface<T> {
|
||||||
setMsgpackImplementation(msgpackProvider)
|
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)
|
// Provider: distance function (resolve BEFORE setupIndex — index uses this.distance)
|
||||||
const nativeDistance = this.pluginRegistry.getProvider<DistanceFunction>('distance')
|
const nativeDistance = this.pluginRegistry.getProvider<DistanceFunction>('distance')
|
||||||
if (nativeDistance) {
|
if (nativeDistance) {
|
||||||
|
|
|
||||||
|
|
@ -117,7 +117,12 @@ export function dequantizeSQ8(quantized: Uint8Array, min: number, max: number):
|
||||||
* @param bMax - Second vector codebook maximum
|
* @param bMax - Second vector codebook maximum
|
||||||
* @returns Approximate cosine distance [0, 2]
|
* @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,
|
a: Uint8Array,
|
||||||
aMin: number,
|
aMin: number,
|
||||||
aMax: number,
|
aMax: number,
|
||||||
|
|
@ -185,6 +190,59 @@ export function distanceSQ8(
|
||||||
return 1 - Math.max(-1, Math.min(1, cosine))
|
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.
|
* Serialize an SQ8 quantized vector to a compact binary format.
|
||||||
*
|
*
|
||||||
|
|
|
||||||
|
|
@ -10,12 +10,14 @@
|
||||||
* - Config defaults: quantization disabled by default (no behavior change)
|
* - 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 { v4 as uuidv4 } from 'uuid'
|
||||||
import {
|
import {
|
||||||
quantizeSQ8,
|
quantizeSQ8,
|
||||||
dequantizeSQ8,
|
dequantizeSQ8,
|
||||||
distanceSQ8,
|
distanceSQ8,
|
||||||
|
distanceSQ8Js,
|
||||||
|
setSQ8DistanceImplementation,
|
||||||
serializeSQ8,
|
serializeSQ8,
|
||||||
deserializeSQ8
|
deserializeSQ8
|
||||||
} from '../../../src/utils/vectorQuantization.js'
|
} from '../../../src/utils/vectorQuantization.js'
|
||||||
|
|
@ -28,6 +30,47 @@ function randomVector(dim: number): number[] {
|
||||||
return Array.from({ length: dim }, () => Math.random() * 2 - 1)
|
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)
|
// Helper: cosine distance for float32 vectors (reference implementation)
|
||||||
function cosineDistance(a: number[], b: number[]): number {
|
function cosineDistance(a: number[], b: number[]): number {
|
||||||
let dot = 0
|
let dot = 0
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue