Commit Graph

30 Commits

Author SHA1 Message Date
Jojii
703cde58a0 Refactor TurboQuantizer (#9095)
* Refactor TurboQuant into common/

# Conflicts:
#	lib/common/turboquant/src/math.rs
#	lib/quantization/src/encoded_vectors_tq.rs
#	lib/quantization/tests/integration/test_tq.rs

* [ai] fix docs

* [ai] tighten function visibilities

* Review remark

* include turboquant in amalgamate

* Rebase

* Move common/turboqunat back (but keep the refactor)

* Move vector_stats.rs back
2026-06-03 14:51:12 +02:00
xzfc
2b5cf80f66 Warn on clippy::wildcard_enum_match_arm (#9096) 2026-05-22 10:46:36 +02:00
Luis Cossío
6598a47933 make quantized_vector_size method static (#9060) 2026-05-22 10:42:50 +02:00
Ivan Pleshkov
d60fb9c77e Apply p square for TQ+ (dont use std+mean) (#8877)
* apply p square

* review remarks and reduce ram

* review remarks

* clean tmp logs

* review remarks

* review remarks

* trigger ci
2026-05-08 13:48:27 +02:00
Jojii
2d075b4998 TQ Hadamard SIMD (#8883)
* [ai+manual] hadamard SIMD

* [ai] Neon

* Minor refactor

* fix import
2026-05-08 13:48:24 +02:00
Ivan Pleshkov
cd2b86491c Tq plus (#8836)
* Add TQ+ ErrorCorrection on top of renorm

Per-coordinate shift+scale fits each rotated, length-rescaled coord onto
the codebook's N(0, 1) grid before quantization. EncodedVectorsTQ::encode
runs a first pass to fit the stats when TQMode::Plus.

Scoring stays correct under renorm's `scaling_factor` framework:
- Asymmetric: precompute_query scales `Q .* D'` and stashes `qm = ⟨Q, M⟩`
  on EncodedQueryTQ; score_precomputed adds qm to raw_dot before applying
  scaling_factor.
- Symmetric: scalar slow path computes `Σ X+_a X+_b D'_i² + xm_a + xm_b
  − ⟨M, M⟩` (xm stored per vector in extras, mm_const cached on
  ErrorCorrection). Result feeds the existing `* v1_scale * v2_scale` arms.
  SIMD reuse for this path is a follow-up.

Storage layout: TQMode::Plus extras are 4 bytes longer (xm appended after
scaling_factor). Zero-vector inputs skip EC application so renorm's
existing zero-norm guard keeps producing score ≈ 0 within tolerance.

VectorStats refactor: streaming `VectorStatsBuilder` so the Plus first
pass can feed Welford with a reused buffer; `build` now takes `dim`
directly and is generic over `T: Into<f64>`.

Integration tests run on both Normal and Plus via rstest cases.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Fix TQ+ recall regression on non-uniform per-coord variance data

Two compounding bugs in the renorm + TQ+ composition:

1. centroid_norm was measured on `X+` centroids, which have chi-squared
   norm distribution across vectors (~10% spread for d=256). renorm
   assumed `cn` should be deterministic `sqrt(d)` (just quantization
   drift), so the per-vector correction ended up amplifying intrinsic
   chi-squared noise into ranking error. Fix: revert EC per coord before
   measuring (`c · D' + M`), matching llama-turbo-quant. The reverted
   centroids approximate `rescaled` which has length `sqrt(d)` exactly
   by construction. dequantize follows the same convention so the
   stored `scaling_factor = l2/cn` round-trips back to the original l2.

2. Asymmetric query path pre-scaled `Q* = R_q · D'` before SIMD encoding.
   The SIMD encoder normalizes by `max(|input|)`, so a query whose coords
   span 5× magnitude (which `R_q · D'` does on real data) loses precision
   on the small-D' coords. Fix: keep `rotated` unscaled, store it as a
   side field on `EncodedQueryTQ`, and use a scalar decode-and-dot path
   for TQ+ (`Σ R_q_i · c_i · D'_i + qm`). SIMD support for this is a
   follow-up.

Catches both via `recall_skewed_data` test on data with 8 spike-variance
input coords. Without the fixes, Bits4 Plus dropped to 0.93 vs Normal
0.98; after the fixes, Plus tracks Normal within 2%.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* TQ+ asymmetric: skip the SIMD encoding entirely

The TQ+ asymmetric path doesn't use the SIMD-encoded query — it goes
through `score_precomputed_ec` with `rotated_query`. So building the
SIMD form was wasted work + memory. Make `data` an `Option` and only
populate it for the cases that actually use it (Normal mode any distance,
TQ+ L1 via dequantize fallback).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Revert TQ+ to SIMD scoring path

The whole point of TQ+ is that scoring code-paths stay identical between
Normal and Plus modes — only the query precomputation changes. We
pre-scale `Q* = R_q · D'` so the existing SIMD raw_dot computes
`⟨Q · D', X+⟩` directly, then add `qm` and apply renorm's scaling_factor.

Drops `score_precomputed_ec` and `EncodedQueryTQ::rotated_query`. The
recall regression that motivated the scalar fallback was entirely from
the `compute_centroid_norm` bug (measuring `‖X+‖` instead of `‖rescaled‖`)
fixed in the prior commit; SIMD precision was a red herring.

`recall_skewed_data` confirms: Bits4 Normal=0.984 / Plus=0.978, Bits2
Normal=0.902 / Plus=0.908.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* TQ+ Bits1: widen query encoding to 12 bits

For 1-bit storage with TQ+, the per-coord `D' = 1/scale` pre-scaling on
the query can push some coords toward the small end of the SIMD encoder's
integer range (which normalizes by `max(|input|)`). At 8 bits those small
coords lose precision; at 12 bits the rounding error drops ~10× per the
existing `test_query_dotprod_matches_reference` parity test.

`Query1bitSimd` is already generic over BITS so this is just a new
`EncodedQueryTQData::Bits1Wide(Query1bitSimd<12>)` variant + a TQ+/Bits1
dispatch in `precompute_query`. Bits2/Bits4 don't need this — their
storage is fine-grained enough that query precision isn't the bottleneck,
and their SIMD encoders aren't generic over BITS today.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* are you happy fmt

* TQ+ Bits1Wide: bump to 16-bit query quantization (kernel max)

12 bits helped on real datasets but not enough — push to the kernel's
ceiling of 16. `Query1bitSimd<BITS>` asserts `BITS ∈ [2, 16]`, so this
is the most precision the existing SIMD path can give us before needing
a wider integer kernel.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* TQ+: shift by per-coord median, not mean

For 1-bit storage the codebook boundary sits at 0, so post-shift values
are quantized purely by sign. Median is the sign-balance point of the
distribution; mean isn't (skewed coords pull mean off the median).

On anisotropic embeddings — dbpedia-openai being the reference case —
mean-based shift produced a ~60/40 biased sign distribution per coord,
losing 1-bit's representational capacity. Median-based shift restores
50/50 and matches llama-turbo-quant's behavior. Higher bit-widths are
less sensitive but still benefit; the codebook boundaries still lie at
distribution-percentile-aware positions when the data is centered on
the median.

Median requires per-coord samples in memory, so cap the stats pass at
10K vectors. Estimates converge fast (~√N) — 10K is plenty even for
million-vector indexes. The encoding pass still processes every vector.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* TQ+ Bits1Wide: revert to 12-bit query quantization

The recall regression on anisotropic data was the mean-vs-median shift,
not query precision. 12 bits is enough headroom for the per-coord D'
pre-scaling and avoids the extra storage of the 16-bit form.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* TQ+ Bits1Wide: bump back to 16-bit query quantization

12 bits helped a bit but not enough on the real dataset. Bump to the
kernel's ceiling. If 16 still isn't enough, the next step is checking
whether the gap is real (re-measure llama branch) before widening the
SIMD integer kernel itself.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* are you happy clippy

* use mean

* 1bit error correction

* review remarks

* are you happy clippy

* review remarks

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-08 13:47:45 +02:00
Ivan Pleshkov
541bb8baf3 TQ SIMD (#8749)
* TQ 4 bit SIMD

* more optimizations

* unroll test

* remove tries

* revert avx 512

* final simd

* use precomputed codebooks

* close to finish

* split to files

* are you happy fmt

* score internal

* 1bit

* 1bit tails

* score_1bit_internal_avx2 tail

* 1bit simd

* 2bit case

* fix 2bit

* fix features

* tails

* fix tests

* less benches

* 1bit tails

* 4bit tails simd

* 64k overflow test

* reuse packing and constants from dev after rebase

* are yoy happy fmt

* fix x64 build

* integration

* symmetric score SIMD

* fix codespell

* better docs

* are you happy clippy

* review remarks

* review remarks
2026-05-08 13:47:36 +02:00
Jojii
361b195f2a TQ Scoring (#8753)
* [AI + Manual] Cosine scoring

* [ai + manual] Add support for dot product

* [ai] re-implement tests including covering both distances (dot+cosine)

* Clarify Cosine/Dot equivalence behavior for pre-TQ quantization storages

* review remarks

---------

Co-authored-by: Ivan Pleshkov <pleshkov.ivan@gmail.com>
2026-05-08 13:47:20 +02:00
Jojii
f48e0cfeee TurboQuant basic Quantization + Encoding (#8692)
* [AI + Manual] TurboQuant quantization

* Code spell

* [Ai + Manual] Integrate TurboQuantizer

* [Ai + Manual] Review tests

* Add Todos for unimplemented parts

* Fix bench

* [manual + ai] Add TqVectorExtras for additional metadata in quantized vectors

* Direct get_centroids implementation
2026-05-08 13:47:17 +02:00
Jojii
522f0ed592 TurboQuant hadamard Rotation (#8657)
* [ai + manual] Implement Hadamard rotation

# Conflicts:
#	lib/quantization/src/turboquant/mod.rs

* [ai + manual] Defer normalization and reduce iterations

* Test improvements + Clippy

* [AI] Don't use fixed-size chunks

* [AI + Manual] Extract permutations and add more tests

* Add normalization + Remove signs2

* Remove signs2 too and improve chunk size function

* [ai] Improve code

* [ai] use reversible LCG for O(1) memory shuffle

* Bench fixes + Remove constant in seed

* [ai] Review Remarks

* [ai] Review remarks

* [ai + manual] In-Place rotations
2026-05-08 13:46:44 +02:00
Luis Cossío
3441e395e4 Chunked vectors with UniversalWrite storage (#8233)
* use CowMultiVector as return type from storages

* add advice to OpenOptions

* Implement ChunkedVectors with generic storage

* rename ChunkedVectors->VolatileChunkedVectors and ChunkedMmapVectors-> ChunkedVectors

* propagate everywhere

fix tests

* [auto] rename BytesRange -> ElementsRange

* [auto] rename BytesOffset -> ElementOffset

* coderabbit nits

---------

Co-authored-by: generall <andrey@vasnetsov.com>
2026-03-26 17:25:24 +01:00
dependabot[bot]
a7739dbf2d build(deps): bump rand_distr from 0.5.1 to 0.6.0 (#8148)
* build(deps): bump rand_distr from 0.5.1 to 0.6.0

Bumps [rand_distr](https://github.com/rust-random/rand_distr) from 0.5.1 to 0.6.0.
- [Release notes](https://github.com/rust-random/rand_distr/releases)
- [Changelog](https://github.com/rust-random/rand_distr/blob/master/CHANGELOG.md)
- [Commits](https://github.com/rust-random/rand_distr/compare/0.5.1...0.6.0)

---
updated-dependencies:
- dependency-name: rand_distr
  dependency-version: 0.6.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>

* Migrate main code base to rand 0.10

* Migrate tests

* Migrate benches

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: timvisee <tim@visee.me>
2026-03-26 17:09:49 +01:00
Ivan Pleshkov
8dc1544916 Use p square to find ranges (#7733)
* use p square to find ranges

* use sample size

* add stopper checks

* review remarks

* review remarks
2026-02-09 23:04:25 +01:00
Ivan Pleshkov
e1c6ec71a9 Scalar quantization encoding parameter (#7602)
* scalar quantization encoding parameter

fix ci

remove method

review remarks

fix ci

* uint8 - int8
2025-12-03 10:23:25 +01:00
Luis Cossío
60b311d601 homogenize bench for aarch64 too (#7628) 2025-12-03 10:20:11 +01:00
Ivan Pleshkov
83bb3a3598 refactor sq score_bytes (#7613) 2025-12-03 10:19:10 +01:00
Ivan Pleshkov
9ec813e0da P-Square one pass quantile estimation method (#7520)
* p square one pass method

* are you happy fmt

* fix n=9 case

* marker struct

* refactor

* proper tests

* add bench

* better namings

* are you happy clippy

* are you happy clippy

* fix additional markers order

* use ArrayVec instead of SmallVec

* use orderer float to sort f64

* Update lib/quantization/src/p_square.rs

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>

* Update lib/quantization/src/p_square.rs

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>

* mispelling

* review remarks

* fix typo

* change float checks order

---------

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>
2025-11-25 11:15:51 +01:00
Ivan Pleshkov
6bc3cc813c EncodedStorage upsert vector (#7048)
* EncodedStorage upsert vector
2025-08-26 13:16:40 +02:00
Ivan Pleshkov
18ba202a76 Remove quantization save functions (#7043)
* remove quantization save

* revert convert_binary_encoding fn

* review remarks

* review remark

* review remarks
2025-08-26 13:16:39 +02:00
Ivan Pleshkov
cdf05b30c3 Wrap quantization Vec<u8> in a separate structure (#7013)
* separate struct for vec u8 quantization storage

* cgf testing for test storage
2025-08-14 14:29:25 +02:00
Ivan Pleshkov
f4ebdd9c9d Appendable quantization storage (#6935)
* appendable qunatization storage

fmt

deprecate count in quantization config

fix arm tests build

qunatized storage vectors count

fix ci

are you happy clippy

fix compat tests

undo rename to `deprecated_count`

remove flusher

remove obsolete functions, don't use hardware counter in ram storage

are you happy clippy

fix gpu build

check ci when revert option

revert last commit

* debug ci

* log offsets

* log offsets

* fix compatibility tests

* deprecate count

* fix arm tests

* fix benches

* Update lib/quantization/src/encoded_vectors_binary.rs

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>

---------

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>
2025-08-11 13:17:02 +02:00
Ivan Pleshkov
fbb92803cc bq scalar query benches (#6831)
* bq scalar query benches

* review remarks
2025-07-17 13:59:04 +02:00
Ivan Pleshkov
bd273ab68f Merge pull request #6728
* # This is a combination of 7 commits.

* SameAsStorage default value

* fix coderabbit warnings

* neon for u8 bq

* sse for u8 bq

* fix windows build

* rename function

* add comments

* fmt

* fix arm build

* review remarks
2025-07-17 13:45:51 +02:00
Ivan Pleshkov
3a916c5d69 Merge pull request #6663
* bq encodings

* are you happy clippy

* are you happy clippy

* are you happy clippy

* are you happy clippy

* gpu tests

* update models

* are you happy fmt

* move additional bits to the end

* fix tests

* Welford's Algorithm

* review remarks

* are you happy clippy

* remove debug println in test

* coderabit nitpicks

* remove unnecessary clone and partialeq

* Use f64 for Welford's Algorithm

* try fix ci

* revert cargo-nextest

* add debug assertions
2025-07-17 13:16:17 +02:00
Tim Visée
8ad2b34265 Bump Rust edition to 2024 (#6042)
* Bump Rust edition to 2024

* gen is a reserved keyword now

* Remove ref mut on references

* Mark extern C as unsafe

* Wrap unsafe function bodies in unsafe block

* Geo hash implements Copy, don't reference but pass by value instead

* Replace secluded self import with parent

* Update execute_cluster_read_operation with new match semantics

* Fix lifetime issue

* Replace map_or with is_none_or

* set_var is unsafe now

* Reformat
2025-03-21 11:38:56 +01:00
Luis Cossío
f110328296 bump and migrate to rand 0.9.0 (#5892)
* bump and migrate to rand 0.9.0

also bump rand_distr to 0.5.0 to match it

* Migrate AVX2 and SSE implementations

* Remove unused thread_rng placeholders

* More random migrations

* Migrate GPU tests

* bump seed

---------

Co-authored-by: timvisee <tim@visee.me>
Co-authored-by: Arnaud Gourlay <arnaud.gourlay@gmail.com>
2025-02-11 14:36:20 +01:00
Jojii
5aee24cc08 Timeout aware hardware counter (#5555)
* Make hardware counting timeout aware

* improve test

* rebuild everything

* fmt

* post-rebase fixes

* upd tests

* fix tests

---------

Co-authored-by: generall <andrey@vasnetsov.com>
2025-01-08 14:03:52 +01:00
xzfc
8bbe66c0fb Reduce debug size (#5556)
* debug size: RawScorerImpl::peek_top_iter

* debug size: TypedMultiDenseVector::multi_vectors() uses

* debug size: stop_condition -> stopped
2024-12-09 11:16:38 +01:00
Jojii
faf645ec2e Hardware counting for quantization (#5369)
* add cpu measurement for quantization

* clippy

* fix tests

* discard hardware counters in benchmarks

* Forwarding hardware counter in distributed setup (#5371)

* make hardware counter available in distributed setup

* clippy

---------

Co-authored-by: generall <andrey@vasnetsov.com>
2024-11-08 11:26:39 +01:00
Ivan Pleshkov
21a3c108bb Move quantization repo (#5096)
* move quantization repo

* are you happy fmt

* are you happy clippy

* remove dumping pq to image

* workspace deps

* are you happy clippy
2024-09-23 16:03:29 +02:00