[diskann-inmem] Add support for spherical quantization by hildebrandmw · Pull Request #1432 · microsoft/DiskANN

Mark Hildebrand (hildebrandmw) · 2026-09-22T17:41:09Z

Add support for spherical quantization in diskann-inmem. As it currently stands, this accepts any diskann_quantization::spherical::iface::Quantizer and has the option for no reranking or reranking using f16 converted full-precision vectors. Reranking can be disabled for index builds to reduce memory consumption.

Search is fixed to use diskann_quantization::spherical::iface::QueryLayout::FullPrecision. This will be made tunable in future changes to diskann-inmem.

This is a large change and is broken into 6 commits of logically related functionality.

Commit 1: Changes to the epoch registry, tags, and slots API.

With the introduction of quantization, we will have multiple Slots implementations being driven by a single Store. This means that:

The reader constructors like Intrusive::reader will no longer work since the parent Store may have a richer Slots generic than just Intrusive.
Reader constructors need to acquire an epoch::Guard externally, meaning that the constructors either need to be unsafe, or we need a way of validating the guard (this PR implements the latter).
When constructing multiple readers from multiple Slots driven by the same Store, it's desirable to have multiple guards. To avoid acquiring multiple guards independently, this means that we need a way of cheaply cloning guards without losing de-registration on drop. This PR implements intrusive reference counting to facilitate this.
Slots implementations may not want to manually mirror the authoritative tags list. In that case, a read-only shared mirror is needed.

For the changes:

diskann-inmem/src/epoch.rs: Adds the reference counting for Guards, using the guard slots in Registry to avoid allocating. In addition, it adds a read only RegistryHandle that can be used to validate that a Guard belongs to a certain Registry.
diskann-inmem/src/tag.rs: Adds an official Authoritative tag list with the ReadOnly variant that can be used by future Slots implementations.
diskann-inmem/src/store/slots.rs: Changes the contract of SlotsConfig to accept both a RegistryHandle and a &Authoritative to the build method. This method now has to be unsafe as the caller assert that the registry handle and tag list are truly the ones being used to manage the Slots.
diskann-inmem/src/store/intrusive.rs: Change the Reader constructor to perform Guard validation.
diskann-inmem/src/store/mod.rs: Update to the new SlotsConfig contract.

The other changes are minor or auxiliary.

Commit 2: More `Slots` implementations

This adds additional Slots implementations needed by quantization.

Simple: A version of Intrusive that is still a dense blob, but uses the tag::ReadOnly introduced in the previous commit as its concurrency control. This will be used to store full-precision vectors for reranking.
Optional: An optional Slots implementation. This will be used to make reranking optional, avoiding the allocation is unused.
Cons: A Slots implementations that combines two other Slots implementations. This will be used to have an Intrusive primary for quantized vectors with an Optional<Simple> for the reranked store.

In addition, this exposes the Simple store under the integration-test flag for stress testing. Infrastructure in diskann-inmem/integration/store/* is adjusted so the write/read pattern validation currently used by the intrusive tests is reused by the simple store.

Commit 3: Factor out shared `ExpandBeam`/`Prune`

This factors out the ExpandBeam and Prune implementations from diskann-inmem/src/repr/full.rs into general structs in diskann-inmem/src/repr/internal/intrusive.rs. To help with this process, new diskann_inmem::repr::{RawDistance, RawQueryDistance} traits are introduced.

The bulk of the changes are in diskann-inmem/src/repr/internal/intrusive.rs and diskann-inmem/src/repr/full.rs.

This PR also changes the ExpandBeam trait to use a Vec<Neighbor<u32>> instead of the old tuple to help with allocation reuse in reranking.

The other changes are largely auxiliary and small - setting up for the next few commits.

Commit 4: Add the `Spherical` representation.

This is where everything comes together. And will be broken down by file.

diskann-inmem/src/repr/internal/simple.rs: A general purpose Reranker based on the Simple slots introduced in commit 2.
diskann-inmem/src/repr/test/mod.rs: This adds a Reference helper struct. When working with Representations - we don't necessarily guarantee the assignment order for internal slots. This can make tests kind of annoying to write. To that end, the Reference track a LogicalId (stable external handle), a SlotId (the internal assigned slot) and the data payload. It is indexable by either LogicalId or SlotId. This is used to test both the reranker and the simple store.
diskann-inmem/src/num.rs: Introduction of the test-only LogicalId and SlotId.
diskann-inmem/src/repr/spherical.rs: The meat and potatoes! Given all the set-up infrastructure work, this is relatively straightforward to add.

The underlying storage is a Cons<Intrusive, Optional<Simple>>. That is, the first slots are an intrusive ones used to store quantized vectors. The second slots are for optional reranking. Sizing, distances, and compression into the first store are all done through an iface::Quantizer trait object from diskann-quantization. Reranking is currently selected through an enum, and can either be disabled (Rerank::None) or use float16 (Rerank::F16).

For reranking, we use the SearchAccessor::post_process, creating one as needed.

Note that spherical search supports multiple different query layouts. Currently, insert uses the SameAsData layout (since distances needed to be compatible with pruning) and search uses the FullPrecision layout. This will be made customizable in the future.

Commit 5: Add integration tests and baselines

This simply hooks up the new Spherical representation to the integration test infrastructure. Note that one baseline is introduced per bit/distance pair. This is mainly to keep multi-threaded tests running more quickly. Reranking is only tested for the 1-bit cases, and only 1-bit/L2 is run for coverage (as doing anything else is unnecessarily slow).

Commit 6: Integrate into `diskann-benchmark`.

What the heading says. Note that so far, performance is roughly on-par with the original in-memory index. Full performance tuning is deferred as future work once more functionality lands.

Codecov Comments Bot (codecov-commenter) · 2026-09-22T21:05:59Z

Codecov Report

❌ Patch coverage is 95.07074% with 108 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.88%. Comparing base (c0d6692) to head (cec8e99).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-benchmark/src/index/inmem2.rs	84.46%	32 Missing ⚠️
diskann-inmem/integration/index/runner.rs	82.88%	19 Missing ⚠️
diskann-inmem/src/repr/spherical.rs	97.47%	17 Missing ⚠️
diskann-inmem/src/store/simple.rs	93.51%	14 Missing ⚠️
diskann-inmem/src/repr/test/mod.rs	90.00%	7 Missing ⚠️
diskann-inmem/integration/store/simple.rs	95.45%	5 Missing ⚠️
diskann-inmem/src/repr/internal/simple.rs	95.55%	4 Missing ⚠️
diskann-inmem/src/store/cons.rs	94.28%	4 Missing ⚠️
diskann-inmem/integration/store/intrusive.rs	89.65%	3 Missing ⚠️
diskann-inmem/src/epoch.rs	99.13%	1 Missing ⚠️
... and 2 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1432      +/-   ##
==========================================
+ Coverage   90.58%   91.88%   +1.30%     
==========================================
  Files         568      579      +11     
  Lines      112411   114556    +2145     
==========================================
+ Hits       101822   105258    +3436     
+ Misses      10589     9298    -1291

Flag	Coverage Δ
miri	`91.88% <95.07%> (+1.30%)`	⬆️
unittests	`91.82% <95.07%> (+1.48%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann-benchmark/src/main.rs	`94.07% <100.00%> (+0.29%)`	⬆️
diskann-inmem/integration/index/object.rs	`99.00% <100.00%> (ø)`
diskann-inmem/integration/main.rs	`81.33% <100.00%> (+0.77%)`	⬆️
diskann-inmem/integration/store/checked.rs	`94.87% <100.00%> (+0.04%)`	⬆️
diskann-inmem/integration/store/mod.rs	`93.36% <100.00%> (+0.02%)`	⬆️
diskann-inmem/integration/support/datatype.rs	`91.62% <100.00%> (-0.07%)`	⬇️
diskann-inmem/src/integration/store/mod.rs	`96.55% <100.00%> (ø)`
diskann-inmem/src/integration/store/simple.rs	`100.00% <100.00%> (ø)`
diskann-inmem/src/num.rs	`98.29% <ø> (ø)`
diskann-inmem/src/provider.rs	`96.70% <100.00%> (+0.81%)`	⬆️
... and 20 more

... and 47 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot

Copilot review overview

🟡 Changes recommended

One or more issues must be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 · 2

Open (3)

Keep inmem2-u8 benchmark registration · New
Use the correct cargo nextest command · New
Fix quantization spelling in error message · New

What changed in this PR

Adds spherical quantization support to diskann-inmem, including optional F16 reranking, new storage backends, EBR/tag infrastructure, integration coverage, and benchmark wiring.

Changes:

Adds Simple, Optional, and Cons slot implementations with shared tag/epoch handling.
Introduces the Spherical representation and reusable search, pruning, and reranking infrastructure.
Adds integration tests, baselines, configuration support, and benchmark registration.

File	Description
diskann-inmem/src/test/mod.rs	Updated as part of this pull request.
diskann-inmem/src/tag.rs	Updated as part of this pull request.
diskann-inmem/src/store/slots.rs	Updated as part of this pull request.
diskann-inmem/src/store/simple.rs	Updated as part of this pull request.
diskann-inmem/src/store/optional.rs	Updated as part of this pull request.
diskann-inmem/src/store/mod.rs	Updated as part of this pull request.
diskann-inmem/src/store/intrusive.rs	Updated as part of this pull request.
diskann-inmem/src/store/cons.rs	Updated as part of this pull request.
diskann-inmem/src/store/checked.rs	Updated as part of this pull request.
diskann-inmem/src/repr/test/mod.rs	Updated as part of this pull request.
diskann-inmem/src/repr/mod.rs	Updated as part of this pull request.
diskann-inmem/src/repr/internal/simple.rs	Updated as part of this pull request.
diskann-inmem/src/repr/internal/mod.rs	Updated as part of this pull request.
diskann-inmem/src/repr/internal/macros.rs	Updated as part of this pull request.
diskann-inmem/src/repr/internal/intrusive.rs	Updated as part of this pull request.
diskann-inmem/src/repr/full.rs	Updated as part of this pull request.
diskann-inmem/src/provider.rs	Updated as part of this pull request.
diskann-inmem/src/num.rs	Updated as part of this pull request.
diskann-inmem/src/lib.rs	Updated as part of this pull request.
diskann-inmem/src/integration/store/simple.rs	Updated as part of this pull request.
diskann-inmem/src/integration/store/mod.rs	Updated as part of this pull request.
diskann-inmem/src/epoch.rs	Updated as part of this pull request.
diskann-inmem/integration/support/datatype.rs	Updated as part of this pull request.
diskann-inmem/integration/store/simple.rs	Updated as part of this pull request.
diskann-inmem/integration/store/mod.rs	Updated as part of this pull request.
diskann-inmem/integration/store/intrusive.rs	Updated as part of this pull request.
diskann-inmem/integration/store/checked.rs	Updated as part of this pull request.
diskann-inmem/integration/main.rs	Updated as part of this pull request.
diskann-inmem/integration/jsons/store-stress.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/store-stress-test.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/integration.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/integration-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/two-bit-l2.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/two-bit-l2-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/two-bit-ip.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/two-bit-ip-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/two-bit-cosine.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/two-bit-cosine-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/one-bit-l2.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/one-bit-l2-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/one-bit-ip.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/one-bit-ip-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/one-bit-cosine.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/one-bit-cosine-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/four-bit-l2.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/four-bit-l2-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/four-bit-ip.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/four-bit-ip-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/four-bit-cosine.json	Updated as part of this pull request.
diskann-inmem/integration/jsons/graph/spherical/four-bit-cosine-baseline.json	Updated as part of this pull request.
diskann-inmem/integration/index/runner.rs	Updated as part of this pull request.
diskann-inmem/integration/index/object.rs	Updated as part of this pull request.
diskann-inmem/DEV.md	Updated as part of this pull request.
diskann-inmem/Cargo.toml	Updated as part of this pull request.
diskann-benchmark/src/index/inmem2.rs	Updated as part of this pull request.
diskann-benchmark/Cargo.toml	Updated as part of this pull request.
Cargo.lock	Updated as part of this pull request.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Change tags to use an authoritative store. Slots can retain a read-only mirror of this store. Update epoch guards to support cheap reference-counted cloning and retain a backedge to their parent Registry. This allows RegistryHandle to validate that a guard belongs to its associated Registry, so Slots can provide safe Reader constructors. Finally, update SlotsConfig to accept both a RegistryHandle and &Authoritative, allowing Slots implementations to retain either. This moves unsafety from reader construction to Slots construction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

- Simple: a dense blob that uses the read-only authoritative tag mirror for concurrency. - Optional: an optional Slots implementation. - Cons: combines two Slots implementations. Additionally, add stress integration tests for the Simple store. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

@hildebrandmw

# DiskANN v0.60.0 ## Major Changes ### `diskann-utils`: Matrix API changes There have been significant changes to the matrix APIs in `diskann-utils`. The previous contents of `diskann_utils::views` have been moved to `diskann_utils::views::rowmajor`. The following type changes have been made | v0.59 | v0.60 | Remarks | |-|-|-| | `views::Matrix<T>` | `views::rowmajor::Owned<T>` | Size changed from 32 bytes to 24 bytes | | `views::MatrixView<'a, T>` | `views::rowmajor::Ref<'a, T>` | Size changed from 32 bytes to 24 bytes | | `views::MatrixViewMut<'a, T>` | `views::rowmajor::Mut<'a, T>` | Size changed from 32 bytes to 24 bytes | | `MatrixBase` inherent methods | `rowmajor::Matrix` and `rowmajor::MatrixMut` trait methods | | The `rowmajor::Matrix` and `rowmajor::MatrixMut` trait methods need to be imported to access most matrix functionality. In addition, the following methods have been renamed and changed: | v0.59 | v0.60 | Remark | |-|-|-| | `row_iter` | `rows` | No longer panics when `ncols == 0` | | `row_iter_mut` | `rows_mut` | No longer panics when `ncols == 0` | | `par_row_iter` | `par_rows` | No longer panics when `ncols == 0` | | `par_row_iter_mut` | `par_rows_mut` | | | `get_row_unchecked` | `row_unchecked` | | | `get_row_unchecked_mut` | `row_unchecked_mut` | | | `as_mut_view` | `as_view_mut` | | | `try_from` constructors | `try_from_data` | | In addition, * The `Index` and `IndexMut` implementations have been removed. Use `Matrix::element` and `Matrix::element_mut` instead. * The `Init` and `Generator` style constructors for `Matrix::new` have been removed. Instead, `rowmajor::Owned::from_fn` and `rowmajor::Owned::from_element` can be used instead. The former's closure has been changed to take an argument `rowmajor::RowCol` indicated the row and column being initialized. Both of these constructors take now take the initializer as the last argument. * Matrix constructors validate against integer overflow and `isize::MAX` allocations via a `rowmajor::Layout<T>` layout. Constructors panic if these validations fail. `try_*` variations of these constructors which return `rowmajor::LayoutError` instead. ### `diskann-utils`: Strided API changes * Structs `StridedBase`, `StridedView`, and `MutStridedView` were collapsed into a single, immutable `Strided<'a, T>`. * `strided::linear_length` was removed, use `strided::Layout::linear_length` instead. * `strided::TryFromErrorLight` was removed. * `strided::TryFromError` changed from a struct to an enum, and `into_inner`/`as_static` were removed. ### `diskann` `InsertStrategy` is no longer a `SearchStrategy` subtrait. Implementations must now provide ```rust type SearchAccessor; type SearchAccessorError; fn insert_search_accessor(...) -> Result<Self::SearchAccessor, Self::SearchAccessorError>; ``` Previously, `insert_search_accessor` had a provided implementation forwarding to `SearchStrategy::search_accessor`. That default is now gone. This change was done to decouple insert element types from search element types. ### `diskann-disk` The publicly re-exported types below were removed in favor of standard buffered I/O: ``` diskann_disk::storage::CachedReader diskann_disk::storage::CachedWriter ``` Consumers should migrate to `std::io::BufReader`/`BufWriter` and the relevant storage-provider reader/writer. `READ_WRITE_BLOCK_SIZE` also changed from `u64` to `usize`. ### `diskann-quantization` These public traits now require `Debug`: ``` spherical::iface::Quantizer spherical::iface::DynQueryComputer spherical::iface::DynDistanceComputer ``` ### `diskann-inmem` (experimental) The old `layers` abstraction was removed and replaced by the `Representation`/`Store`/`Slots` architecture. And support for a spherically quantized backed was added. # Full Change Log * [diskann-inmem] Prepare code for quantization and beyond by @hildebrandmw in #1352 * Add Neon inner product U8xU1, U8xU2 and U8xU4 kernels by @pfoxARM in #1372 * [CI] Check all workspace features and private docs by @wuw92 in #1400 * [CI] Check Rust license headers by @wuw92 in #1401 * [disk] Replace custom cached I/O with standard buffered readers and writers by @wuw92 in #1396 * [dikkann_wide] Integer SIMDMinMax by @suri-kumkaran in #1407 * Add Neon Squared L2 kernels for USlice2 and USlice4 by @pfoxARM in #1398 * Bump the github-actions group with 3 updates by @dependabot[bot] in #1406 * [CI] Check workspace direct dependency versions by @wuw92 in #1410 * Train fresh disk-search PQ codebooks for each index build by @xwj-ox in #1299 * [diskann-quantization] Make most things in `diskann-quantization/spherical/iface.rs` debuggable by @hildebrandmw in #1414 * [diskann] Make `InsertStrategy` independent of `SearchStrategy` by @hildebrandmw in #1393 * Use #[expect] to detect unnecessary lint suppressions by @wuw92 in #1409 * [diskann-garnet] Store attributes during insert, not after by @metajack in #1412 * [diskann-wide] Fix `alias!` by @hildebrandmw in #1417 * [diskann-utils] Simplify `StridedView`. by @hildebrandmw in #1376 * [diskann-wide] More Neon Zip/Unzip impls by @hildebrandmw in #1405 * [diskann-garnet] Deserialize max id for quantization on load by @metajack in #1419 * More ergonomic `spherical::iface::Quantizer` constructors by @hildebrandmw in #1418 * [diskann-bftree] Clamp number of concurrency test parallel workers by @hildebrandmw in #1433 * Tweak behavior of adaptive L by @magdalendobson in #1354 * Bump the github-actions group with 2 updates by @dependabot[bot] in #1436 * Reform Range Search by @magdalendobson in #1337 * PiPNN 2/6: add numerical kernels by @SeliMeli in #1287 * Matrix overhaul 1 of N - Data Accessors by @hildebrandmw in #1415 * Add `assert_contains!` helper by @hildebrandmw in #1421 * Matrix overhaul 2 of N - `Matrix` invariants by @hildebrandmw in #1420 * [CI] Run SDE tests for diskann with all features by @SeliMeli in #1440 * [diskann-wide] Add Neon `vabdq_u8` and `vabdq_s8`. by @hildebrandmw in #1437 * [diskann-inmem] Add support for spherical quantization by @hildebrandmw in #1432 * [benchmark-runner] Bump edition to 2024 by @hildebrandmw in #1444 * [diskann-wide] Add byte-vector reinterpretations for u32x8 and u64x8 by @partychen in #1445 * Propagate invalid matrix layout errors in production paths by @arrayka in #1453 * Add sparse vector distance kernels (f32/f16 L2, inner product, cosine) by @mana-agarwal in #1434 * Bump the github-actions group with 3 updates by @dependabot[bot] in #1454 * [diskann-utils] Matrix Update 3 of N (constructors) by @hildebrandmw in #1451 * [diskann-utils] Matrix Update 4 of 4 by @hildebrandmw in #1457 ## New Contributors * @mana-agarwal made their first contribution in #1434 **Full Changelog**: v0.59.0...v0.60.0

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-spherical branch from b4719f0 to 326d152 Compare September 22, 2026 21:36

Mark Hildebrand (hildebrandmw) marked this pull request as ready for review September 22, 2026 22:10

Mark Hildebrand (hildebrandmw) requested review from a team and a lite review from Copilot September 22, 2026 22:10

Copilot started reviewing on behalf of Mark Hildebrand (hildebrandmw) September 22, 2026 22:12

View session

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-spherical branch from 326d152 to a6340a0 Compare September 22, 2026 22:14

Copilot AI reviewed Sep 22, 2026

View reviewed changes

Comment thread diskann-benchmark/src/index/inmem2.rs Outdated

Comment thread diskann-inmem/DEV.md Outdated

Comment thread diskann-inmem/integration/index/runner.rs Outdated