Repository navigation
[diskann-inmem] Add support for Product Quantization - #1469
Open
Mark Hildebrand (hildebrandmw) wants to merge 7 commits into
Open
Mark Hildebrand (hildebrandmw) wants to merge 7 commits into
Mark Hildebrand (hildebrandmw) wants to merge 7 commits into
Conversation
Mark Hildebrand (hildebrandmw)
added this pull request to stack #1468
October 7, 2026 00:03
Mark Hildebrand (hildebrandmw)
force-pushed
the
mhildebr/inmem-pq-2
branch
from
October 9, 2026 15:43
0a45cfa to
f9ddb82
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1469 +/- ##
==========================================
+ Coverage 90.94% 92.03% +1.08%
==========================================
Files 586 590 +4
Lines 118391 119785 +1394
==========================================
+ Hits 107675 110243 +2568
+ Misses 10716 9542 -1174
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Mark Hildebrand (hildebrandmw)
marked this pull request as ready for review
October 9, 2026 16:52
Mark Hildebrand (hildebrandmw)
requested review from
a team
and
a balanced review from Copilot
October 9, 2026 16:52
Contributor
Author
Benchmarking JSON{
"search_directories": [
"redacted"
],
"output_directory": null,
"jobs": [
{
"type": "inmem2",
"content": {
"build": {
"alpha": 1.2,
"l_build": 50,
"max_degree": 42,
"num_threads": 8,
"pruned_degree": 32
},
"data": {
"data": "wikipedia/wikipedia_base_100K.bin",
"data_type": "float32",
"distance": "inner_product"
},
"search": {
"groundtruth": "wikipedia/wikipedia-100k.gt.bin",
"num_threads": [
8
],
"queries": "wikipedia/wikipedia_query.bin",
"reps": 3,
"runs": [
{
"recall_k": 10,
"search_l": [
10,
20,
30,
40,
50,
60,
70,
80,
90,
100
],
"search_n": 10
}
]
},
"quantization": {
"product": {
"chunks": 192,
"centers": 256,
"seed": 1
}
}
}
},
{
"type": "inmem2",
"content": {
"build": {
"alpha": 1.2,
"l_build": 50,
"max_degree": 42,
"num_threads": 8,
"pruned_degree": 32
},
"data": {
"data": "msturing/msturing-base-100k.fbin",
"data_type": "float32",
"distance": "squared_l2"
},
"search": {
"groundtruth": "msturing/msturing-100k.gt.10k.bin",
"num_threads": [
8
],
"queries": "msturing/queries-10k.fbin",
"reps": 3,
"runs": [
{
"recall_k": 10,
"search_l": [
10,
20,
30,
40,
50,
60,
70,
80,
90,
100
],
"search_n": 10
}
]
},
"quantization": {
"product": {
"chunks": 25,
"centers": 256,
"seed": 1
}
}
}
},
{
"type": "inmem2",
"content": {
"build": {
"alpha": 1.2,
"l_build": 50,
"max_degree": 42,
"num_threads": 8,
"pruned_degree": 32
},
"data": {
"data": "openai-v3/base-100k.bin",
"data_type": "float32",
"distance": "cosine"
},
"search": {
"groundtruth": "openai-v3/gt-100k.bin",
"num_threads": [
8
],
"queries": "openai-v3/query.bin",
"reps": 3,
"runs": [
{
"recall_k": 10,
"search_l": [
10,
20,
30,
40,
50,
60,
70,
80,
90,
100
],
"search_n": 10
}
]
},
"quantization": {
"product": {
"chunks": 384,
"centers": 256,
"seed": 1
}
}
}
},
{
"type": "graph-index-build-pq",
"content": {
"index_operation": {
"search_phase": {
"groundtruth": "wikipedia/wikipedia-100k.gt.bin",
"num_threads": [
8
],
"queries": "wikipedia/wikipedia_query.bin",
"reps": 3,
"runs": [
{
"recall_k": 10,
"search_l": [
10,
20,
30,
40,
50,
60,
70,
80,
90,
100
],
"search_n": 10
}
],
"search-type": "topk"
},
"source": {
"alpha": 1.2000000476837158,
"backedge_ratio": 1.0,
"data": "wikipedia/wikipedia_base_100K.bin",
"data_type": "float32",
"distance": "inner_product",
"index-source": "Build",
"insert_retry": null,
"l_build": 50,
"max_degree": 32,
"multi_insert": null,
"num_threads": 8,
"save_path": null,
"start_point_strategy": "medoid"
}
},
"max_fp_vecs_per_prune": 0,
"num_pq_chunks": 196,
"seed": 13076402859301299683,
"use_fp_for_search": false
}
},
{
"type": "graph-index-build-pq",
"content": {
"index_operation": {
"search_phase": {
"groundtruth": "msturing/msturing-100k.gt.10k.bin",
"num_threads": [
8
],
"queries": "msturing/queries-10k.fbin",
"reps": 3,
"runs": [
{
"recall_k": 10,
"search_l": [
10,
20,
30,
40,
50,
60,
70,
80,
90,
100
],
"search_n": 10
}
],
"search-type": "topk"
},
"source": {
"alpha": 1.2000000476837158,
"backedge_ratio": 1.0,
"data": "msturing/msturing-base-100k.fbin",
"data_type": "float32",
"distance": "squared_l2",
"index-source": "Build",
"insert_retry": null,
"l_build": 50,
"max_degree": 32,
"multi_insert": null,
"num_threads": 8,
"save_path": null,
"start_point_strategy": "medoid"
}
},
"max_fp_vecs_per_prune": 0,
"num_pq_chunks": 25,
"seed": 13076402859301299683,
"use_fp_for_search": false
}
},
{
"type": "graph-index-build-pq",
"content": {
"index_operation": {
"search_phase": {
"groundtruth": "openai-v3/gt-100k.bin",
"num_threads": [
8
],
"queries": "openai-v3/query.bin",
"reps": 3,
"runs": [
{
"recall_k": 10,
"search_l": [
10,
20,
30,
40,
50,
60,
70,
80,
90,
100
],
"search_n": 10
}
],
"search-type": "topk"
},
"source": {
"alpha": 1.2000000476837158,
"backedge_ratio": 1.0,
"data": "openai-v3/base-100k.bin",
"data_type": "float32",
"distance": "cosine",
"index-source": "Build",
"insert_retry": null,
"l_build": 50,
"max_degree": 32,
"multi_insert": null,
"num_threads": 8,
"save_path": null,
"start_point_strategy": "medoid"
}
},
"max_fp_vecs_per_prune": 0,
"num_pq_chunks": 384,
"seed": 13076402859301299683,
"use_fp_for_search": false
}
}
]
}
|
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
A duplicated match arm can fail warnings-as-errors Clippy checks, alongside several documentation and diagnostic issues.
5 open findings
What changed in this PR
Adds Product Quantization support to diskann-inmem, including shared reranking infrastructure, integration coverage, and benchmarking.
Changes:
- Implements PQ search, insertion, pruning, metrics, reranking, and distance-table pooling.
- Refactors spherical quantization to share internal metric and reranking logic.
- Adds integration baselines and benchmark support.
| File | Description |
|---|---|
diskann-utils/src/object_pool.rs |
Adds empty capacity-bounded pool construction. |
diskann-inmem/src/repr/test/mod.rs |
Shares quantized representation test helpers. |
diskann-inmem/src/repr/spherical.rs |
Adopts shared quantization internals. |
diskann-inmem/src/repr/product.rs |
Implements the PQ representation. |
diskann-inmem/src/repr/mod.rs |
Exports PQ APIs. |
diskann-inmem/src/repr/internal/quantization/rerank.rs |
Adds shared reranking logic. |
diskann-inmem/src/repr/internal/quantization/mod.rs |
Adds internal metric conversions. |
diskann-inmem/src/repr/internal/mod.rs |
Registers quantization internals. |
diskann-inmem/integration/main.rs |
Registers PQ regression tests. |
diskann-inmem/integration/jsons/graph/product/l2.json |
Configures L2 PQ tests. |
diskann-inmem/integration/jsons/graph/product/l2-baseline.json |
Records L2 baselines. |
diskann-inmem/integration/jsons/graph/product/ip.json |
Configures inner-product PQ tests. |
diskann-inmem/integration/jsons/graph/product/ip-baseline.json |
Records inner-product baselines. |
diskann-inmem/integration/jsons/graph/product/cosine.json |
Configures cosine PQ tests. |
diskann-inmem/integration/jsons/graph/product/cosine-baseline.json |
Records cosine baselines. |
diskann-inmem/integration/index/runner.rs |
Builds PQ integration indexes. |
diskann-inmem/integration/index/object.rs |
Enables PQ index dispatch. |
diskann-benchmark/src/index/inmem2.rs |
Adds PQ benchmark training and execution. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Mark Hildebrand (hildebrandmw)
force-pushed
the
mhildebr/inmem-pq-2
branch
2 times, most recently
from
October 9, 2026 20:01
514ba3a to
4711c41
Compare
Mark Hildebrand (hildebrandmw)
marked this pull request as draft
October 12, 2026 22:42
Mark Hildebrand (hildebrandmw)
marked this pull request as ready for review
October 12, 2026 22:42
Mark Hildebrand (hildebrandmw)
force-pushed
the
mhildebr/inmem-pq-2
branch
from
October 13, 2026 20:43
4711c41 to
e153d29
Compare
Mark Hildebrand (hildebrandmw)
force-pushed
the
mhildebr/inmem-pq-2
branch
from
October 13, 2026 21:07
e153d29 to
f110aec
Compare
Harsha Simhadri (harsha-simhadri)
approved these changes
Oct 14, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add support for PQ for
diskann-inmem. This follows the work in #1458 and #1432.For reviewers, this PR is broken into 6 commits:
Commit 1
Adds a
with_capacityconstructor toObjectPoolto create an initially empty pool. This is used for distance table buffering in the PQ repr.Commit 2
Factors out
Rerankertorepr/internal/quantization/rerank.rs. This will be shared between spherical and PQ. In addition, an internalMetricenum is introduced for the three currently supported metrics:SquaredL2,InnerProduct, andCosine.I kept
spherical::Rerankas a public version of the internalRerankenum to try to decouple the public configuration API from how things are done on the back end. Yes, this is a little bit of code duplication, but keeps the API looking likeThe amount of duplication is small, keeps the interface consistent, and lets us experiment on the back end a little.
Additionally, the routines for
test_expand_beamandtest_pruneare moved out of thesphericaltest module and intorepr/test/mod.rs. These will be reused by the PQ representation.Commit 3
Adds a generic
Quantizedrepresentation. The extension point is a [Quantizer] trait which contains methods for quantizing data and creating distance computers. Distance computation uses BYOTE to allow implementations to create multiple specialized computers as they see fit. These specialized computers are inlined directly into theExpandBeamandPrunetrait objects used for search and prune respectively.The goal behind this approach is to make it easy to implement scalar, minmax, and binary quantizers (since we just need to implement the
Quantizertrait) and to make development of new quantizers easy. And hopefully BYOTE means that the performance cost of this is minimal.One final advantage is that the
Quantizedrepresentation is a single concrete type, minimizing monomorphization bloat of this approach.Commit 4
This actually adds the PQ
Quantizer.Query distances are done using the usual trick of pre-populating a lookup table with distances between query chunk fragments and the PQ pivots (see #1458 for the implementation). Since these tables are fairly large (~200KB for a 192 PQ chunks), an
ObjectPoolis used to reuse allocations. The configuration for this pool inConfigis such that an unboundedObjectPoolis not allowed.Quant-quant distances use the
PaddedTableintroduced in #1458.Tests are similar to spherical, but significantly easier since PQ can losslessly represent the test dataset.
Commit 5
Adds integration test support for PQ and baselines covering L2, inner-product, and cosine.
Commit 6
Adds support in
diskann-benchmarkfor the new PQ backend.Benchmarks
The infrastructure in #1458 pays off. This is considerably faster than the PQ index in
diskann-providers. A benchmarking input JSON is included in the comment - highlights are reported here.Note that the performance increase is somewhat exaggerated due to the small size of the dataset (less memory bandwidth because more fits in cache), which highlights perf differences that a bandwidth-constrained environment would hide.