Repository navigation
Conversation
There was a problem hiding this comment.
🔵 Needs a closer look
The new unsafe architecture-specific SIMD kernel and broad dispatch integration warrant final human validation.
0 open findings
What changed in this PR
Adds end-to-end SphericalQ 3-bit support, including optimized AVX2 distance computation.
Changes:
- Adds 3-bit compression, dispatch, serialization, and provider support.
- Implements a bounded-tail AVX2 inner-product kernel.
- Registers 3-bit benchmarks and expands correctness coverage.
| File | Description |
|---|---|
diskann-quantization/src/spherical/vectors.rs |
Adds 3-bit distance tests and metadata-aware references. |
diskann-quantization/src/spherical/quantizer.rs |
Enables 3-bit compression. |
diskann-quantization/src/spherical/iface.rs |
Adds dispatch, plans, and serialization support. |
diskann-quantization/src/spherical/__codegen/x86_64.rs |
Adds 3-bit codegen inspection entry points. |
diskann-quantization/src/bits/distances.rs |
Implements and tests the AVX2 kernel. |
diskann-providers/src/model/graph/provider/async_/inmem/spherical.rs |
Enables 3-bit in-memory stores. |
diskann-benchmark/src/inputs/exhaustive.rs |
Accepts compatible 3-bit benchmark queries. |
diskann-benchmark/src/index/inmem/spherical.rs |
Registers 3-bit graph benchmarks. |
diskann-benchmark/src/exhaustive/spherical.rs |
Registers 3-bit exhaustive benchmarks. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is ❌ Your patch status has failed because the patch coverage (85.93%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #1474 +/- ##
=========================================
Coverage 92.07% 92.07%
=========================================
Files 586 591 +5
Lines 119997 122384 +2387
=========================================
+ Hits 110488 112690 +2202
- Misses 9509 9694 +185
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
|
Just did a quick glance through of the perf numbers - something seems off. This is saying that the index build times for PQ are comparable to spherical? This definitely should not be the case - spherical builds should be significantly faster. |
Yes SphericalQ quantizer training is significantly faster than PQ. The earlier “build time” included graph construction, which obscured this difference. Wikipedia-Cohere, 1M × 768; same 50K training sample, 16 encoding threads, median of three runs:
PQ’s encoder is currently better optimized: it uses transposed codebooks and batched SIMD, while SphericalQ3 encodes each vector separately and searches for its scaling factor. I'll try to improve it, it does seem weird that SphericalQ has same or even worse build time than PQ. |
4ce63c7 to
4f3fa48
Compare
Implement packed 3-bit inner products and enable spherical quantization, serialization, and existing provider/benchmark paths. Reuse shared distance tests across bit widths and architectures, including long-vector tails and stored-metadata references. Add L2, inner-product, and cosine codegen entry points without new dependencies.
4f3fa48 to
19eab8a
Compare
Mark Hildebrand (hildebrandmw)
left a comment
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks - there's some clean-up work to do, particularly around serialization. If you don't need serialization for this current line of work, can we leave it off?
Otherwise, we'll need to refactor deserialization at least to allow callers to control which bit widths they wish to support.
| // Full // | ||
| ////////// | ||
|
|
||
| fn test_full_distance_tails<const NBITS: usize>( |
There was a problem hiding this comment.
What is this test covering that isn't already covered in test_full-distance?
There was a problem hiding this comment.
The additional coverage was explicit misalignment, longer vectors, and non-integer query values. Folded those cases into the existing test_full_distance helper and removed the separate tails helper.
There was a problem hiding this comment.
I don't think this additional coverage is needed.
- Miri will already allocated base byte vectors at strange alignments - and Miri is really the only situation that would catch unannotated unaligned loads since
x86andaarch64are fine with unaligned loads and so will happen to work. - What additional coverage is given by longer vectors? Longer runs take significant time under Miri. We need only exercise enough of the implementation that the various unrolling bounds are hit. If the current upper bound is insufficient, then we should use more precise, architecture dependent upper bounds like raw bit distances.
- Also, what additional coverage is provided by using non-integer values for these tests?
There was a problem hiding this comment.
Thank you! So
- alignment checking is left to Miri.
- Longer vectors make no sense since current dimension bounds already cover 32/16/8-element boundaries and tail paths.
- Non-integer values can be used to verify numerical precision. However, since there are no explicit precision requirements or specific issues to reproduce, there is no need to have the additional coverage.
Removed additional coverage.
Inline the V3 kernel, reuse architecture dispatch, and consolidate distance tests. Downcast 3-bit V4 computers to V3 and preserve the generic deserializer support set. Restore independent reference formulas with targeted 3-bit rounding tolerances. Validation: quantization tests with flatbuffers/codegen, spherical provider tests, baseline distance tests, cargo fmt --all, and workspace Clippy with warnings denied.
Reference Issues/PRs
What does this implement/fix? Briefly explain your changes.
Adds SphericalQ3 support and an AVX2 inner-product kernel for full-precision queries against densely packed 3-bit codes. The kernel uses four accumulators for all dimensions, with bounded tail loads and no padding requirement.
The change includes compression, distance dispatch, serialization, existing in-memory provider and benchmark registration, shared correctness tests, and codegen inspection entry points.
The motivation is to provide another memory/accuracy tradeoff. Compared with PQ using two dimensions per chunk, SphericalQ3 achieves similar recall and broadly comparable end-to-end latency while reducing query-process memory by approximately 11–17% in the experiments below. Recall differences range from −0.740 to +0.028 percentage points.
Any other comments?