In-Mem 2.0

[hildebrandmw]

Introduce a second in-memory provider with the intention of replacing the current provider.

Why

The RFC outlines much of the motivation. In short, the goal here is to:

• Make the provider safe under concurrent inserts/search/deletes.
• Support proper external/internal ID translation.
• Improve test coverage
• Do so with minimal performance overhead.

The concurrency argument comes from the epoch-based-reclamation (EBR) protection scheme for internal slots. See the RFC for more details.

Known Follow-up Items

• Perf Parity: This generally has pretty good performance, but could use a little more tuning to bring it fully on-par with our current inmem index.
• Quantization: This initial prototype is lacking quantization support. Adding quantization is relatively straightforward in the primary Store, but will need some thought on how to add the reranking layer. This shouldn't be an architectural blocker, though.
• Hybrid PQ: A larger open question is how to support the max_fp_vecs feature of our current PQ implementation, which reads some full-precision and some quantized vectors during prune. Like with quantization in general, I think this is not a fundamental issue.
• Support for non-uniform sized items in slots: For this, I'm mainly thinking of multi-vectors where the number of vectors within each multi-vector can vary. In the context of multi-vectors, fast element access is less important than for traditional vectors as distance computations take considerably longer.

Suggested Reviewing Order

The majority of this PR is in a new diskann-inmem crate. To facilitate testing, this crate has an "integration-test" feature, which enables the code in diskann-inmem/src/integration. This code is an unstable public reexport of internal types meant only for consumption in the diskann-inmem/integration integration test binary.

The integration test binary is powered by diskann-benchmark-runner.

diskann-inmem

Independent Low-Level Utilities

• num.rs: Strong type utilities for byte and alignment representations.
• buffer.rs: A miri-compliant version of AlignedMemoryVectorStore. This type allows vectors/neighbors to be stored in a single larger allocation. The use of RawSlice allows slots within the Buffer to be inspected and manipulated without forcing reference materialization (which is important to prevent aliasing).
• neighbors.rs: The new version of SimpleNeighborVectorProviderAsync. This reuses the sharded-lock idea, but provides additional utility, including the ability to perform read-modify-write operations on adjacency lists.
• counters.rs: Event counters. When the "integration-test" feature is not enabled, counters become a no-op. These are enabled for testing to monitor changes.
• sharded.rs: An external-to-internal ID translation utility. The main trick with this struct is to provide utilities like Sharded::occupied_entry, which locks and returns an external/internal mapping. The proxy Entry struct is important as it verifies that such a mapping exists and provides an infallible way of deleting the mapping. This is chained with higher level operations (e.g. Provider::delete) to delete both the ID-mapping and the internal data-slot in lock-step.

Concurrency Protocol

The concurrency protocol is built upon three main layers:

• tag.rs: An atomic slot tag for controlling access to data.
• epoch.rs: The central registry where readers register and deregister. This is the crux of this PR and probably the most important file.
• store.rs: A binary blob store built on top of epoch.rs to provide the safe concurrent store for data. This provides the following operations:
  • Storage of binary data in "slots".
  • Reading of data in slots (provided by Reader).
  • Tracking on the valid/invalid state slots.
  • Finding available slots into which new data can be inserted.
  • Safe retirement of slots and eventual reclamation.

The Store in store.rs has some help from freelist.rs to accelerate locating available slots internally.

Testing: epoch.rs has unit tests with injectable delays to set up known pathological orderings. The sequencing is helped by test/sequencer.rs. In addition, test/epoch.rs includes a direct stress test for the Registry. This is particularly helpful when run under Miri, which has the ability to detect some race conditions.

A larger concurrency stress test lives in the integration-test binary. This directly tests store.rs by spinning up readers, writers, and retirers and hammers a single Store. Data is read and written into the store in a knowable pattern, allowing readers to detect torn reads, implying a race condition. For this PR, I ran the following stress test file

{
  "search_directories": [
  ],
  "output_directory": null,
  "jobs": [
    {
      "type": "store-stress",
      "content": {
        "capacity": 8192,
        "duration_secs": 600,
        "entry_bytes": 256,
        "low_watermark": 4096,
        "max_ops": 50000000000,
        "readers": 32,
        "retirers": 16,
        "seed": 11935966405698895599,
        "writers": 32
      }
    }
  ]
}

with the command

cargo run --package diskann-inmem \
  --bin integration-test \
  --features integration-test \
  --release -- \
   run --input-file stress.json --output-file temp.json

The generated output was

readers:       32
writers:       32
retirers:      16
capacity:      8192
entry_bytes:   256
low_watermark: 4096
duration_secs: 600
max_ops:       50000000000
seed:          11935966405698895599
elapsed_secs:  600.001182962
reads:         110206469888
acquires_ok:   248086767
acquires_fail: 619145
retires_ok:    248082573
retires_fail:  243312644
reclaims:      108847566
transitions:   42006520
peak_live:     8131

While not a proof of correctness, this is a pretty decent stress test.

Providers

The implementation of the data provider is split into two logical pieces. The first lives in layers/ and is focused on computing distances. With this approach, I am trying to avoid the need to replicate Accessors and Strategys for each future quantization type. layers/full.rs is the full-precision implementation. One thing to note is the use of the FullPrecision marker trait from which the implementations of the whole layers API is derived for layers::Full. This allows users to include just a T: FullPrecision trait bound and really simplifies the generics upstream.

Within provider.rs - my goal here is to minimize the use of generics as much as possible. In particular for search, I use a trait object for ExpandBeam. When coupled with the layers::QueryDistance API, we can create implementations where the distance function is inlined and the number of prefetch instructions can be tailored to the data length. Fully optimizing this is still a work-in-progress.

Another thing to call out in provider.rs is the care needed for data insertion and deletion. Since external/internal ID translation is supported, we need to ensure that the translation table stays in-sync with the internal store. On insert, if we allocate an internal slot only to find the external ID already exists, we need to abort the operation rather than publish the internal slot. Similarly on delete, we first need to establish if the external/internal ID mapping exists. If so, then we can try to retire the slot. If slot retiring fails (it shouldn't, but bugs can happen) - we need to not commit the ID mapping deletion and instead return an error.

Testing

Like the concurrency stress tests, testing uses the integration-test binary. Here, the 10k YFCC dataset is used for non-trivial runs. Using the A/B functionality in diskann-benchmark-runner, we can compare against checked-in baselines. This allows us to capture rich metrics for recall, number of operations, etc. and update easily. The main logic for checking and reporting baseline mismatches is in integration/support/check.rs. The goal is to summarize all such mismatches for presentation to provide the highest signal possible.

diskann-benchmark

Integration into the benchmarks is straightforward. I elected to put everything in a single file to minimize disruption. I'm trying an approach of using diskann_benchmark_runner::Input::from_raw to separate out the deserialization types from the actual inputs, allowing richer types (e.g., a full diskann_benchmark_core::streaming::bigann::RunBook) to be loaded.

Also note how relatively simple the streaming benchmark integration is. Since the new inmem provider supports ID translation and internal slot allocation, it does not need the same level of hand-holding that the current inmem provider needs.

[Copilot]

Pull request overview

This PR introduces “In-Mem 2.0” by adding a new diskann-inmem crate implementing a concurrent-safe in-memory provider (EBR + slot tags + blob store), along with integration/stress testing and benchmark-runner plumbing to exercise it and compare against checked-in baselines.

Changes:

• Add new diskann-inmem crate (epoch registry, slot tags, store, ID translation, distance layers) plus unit + integration stress tests and baseline regression harness.
• Integrate inmem2 into diskann-benchmark behind a new inmem2 feature and extend diskann-benchmark-runner formatting/UX helpers used by the integration binary.
• Wire the new crate into the workspace (Cargo.toml, Cargo.lock) and CI feature matrix.

Reviewed changes

Copilot reviewed 48 out of 49 changed files in this pull request and generated 16 comments.

Show a summary per file

File	Description
rfcs/01206-inmem2.md	New RFC documenting the concurrency motivation and EBR/tag protocol design.
diskann-vector/src/distance/implementations.rs	Adjusts SIMD dispatch hook to use run2_inline.
diskann-inmem/src/test/sequencer.rs	Adds a test sequencing helper (condvar + state machine).
diskann-inmem/src/test/mod.rs	Introduces diskann-inmem unit-test module structure.
diskann-inmem/src/test/epoch.rs	Adds directed stress tests for the epoch Registry.
diskann-inmem/src/tag.rs	Implements atomic slot tags + protocol documentation and tests.
diskann-inmem/src/sharded.rs	Adds sharded external↔internal ID mapping utility with tests.
diskann-inmem/src/num.rs	Adds Bytes/Align strong types with tests.
diskann-inmem/src/lib.rs	New crate root wiring modules, exports, and internal helper macro.
diskann-inmem/src/layers/mod.rs	Introduces distance-layer traits (Layer, Search, QueryDistance, etc.).
diskann-inmem/src/layers/full.rs	Implements full-precision distance layer + specialization and tests.
diskann-inmem/src/integration/store.rs	Public (integration-test-only) wrapper around internal store::Store.
diskann-inmem/src/integration/mod.rs	Integration-test-only module exports.
diskann-inmem/src/integration/counters.rs	Integration-test-only counter snapshot type for reporting.
diskann-inmem/src/freelist.rs	Adds a freelist helper for finding/recycling internal slot IDs (with tests).
diskann-inmem/src/counters.rs	Adds counters/no-op counters depending on integration-test feature.
diskann-inmem/src/buffer.rs	Adds raw buffer + RawSlice abstraction for miri-friendly storage (with tests).
diskann-inmem/integration/support/tolerance.rs	Defines an “empty tolerance” input type for runner regression checks.
diskann-inmem/integration/support/mod.rs	Integration support module wiring.
diskann-inmem/integration/support/io.rs	Adds dataset loading + conversion helper for integration tests.
diskann-inmem/integration/support/check.rs	Adds structured baseline mismatch aggregation/printing utilities.
diskann-inmem/integration/store.rs	Adds store concurrency stress benchmark binary implementation.
diskann-inmem/integration/main.rs	Adds integration-test runner binary entrypoint and regression harness.
diskann-inmem/integration/jsons/store-stress.json	Adds example stress-test job JSON.
diskann-inmem/integration/jsons/store-stress-test.json	Adds a smaller stress-test job JSON for CI/unit testing.
diskann-inmem/integration/jsons/integration.json	Adds integration-test job definitions (YFCC 10k A/B runs).
diskann-inmem/integration/jsons/integration-baseline.json	Adds checked-in baseline outputs for regression comparison.
diskann-inmem/integration/jsons/checks.json	Adds checker config using empty tolerances.
diskann-inmem/integration/index/tests.rs	Adds index-level insert + knn execution helpers and result DTOs.
diskann-inmem/integration/index/runner.rs	Adds runner glue: deserialize inputs, build provider, run regression checks.
diskann-inmem/integration/index/object.rs	Adds trait-object interface for an index + counter/metric DTOs.
diskann-inmem/integration/index/mod.rs	Wires index integration modules and registration.
diskann-inmem/DEV.md	Documents how to run full test suite with integration-test enabled.
diskann-inmem/Cargo.toml	New crate manifest, deps, bin target, and integration-test feature.
diskann-inmem/.clippy.toml	Clippy config allowing unwrap/expect/panic in tests for this crate.
diskann-benchmark/src/index/mod.rs	Registers new inmem2 benchmarks behind feature flag.
diskann-benchmark/Cargo.toml	Adds optional dependency and feature flag for diskann-inmem.
diskann-benchmark-runner/src/utils/fmt.rs	Adds KeyValue formatter utility and tests.
diskann-benchmark-runner/src/files.rs	Implements Display for InputFile to improve formatting.
diskann-benchmark-runner/src/app.rs	Minor formatting fix for mismatch reporting output.
Cargo.toml	Adds diskann-inmem to workspace + dependency table; adds profile.samply.
Cargo.lock	Locks new diskann-inmem package and dependencies into the workspace lockfile.
.github/workflows/ci.yml	Enables integration-test and inmem2 in CI feature set.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.31981% with 150 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.82%. Comparing base (ae25484) to head (a4ee3b2).

Files with missing lines	Patch %	Lines
diskann-inmem/src/provider.rs	88.02%	69 Missing ⚠️
diskann-inmem/src/store.rs	96.08%	16 Missing ⚠️
diskann-inmem/src/epoch.rs	96.01%	13 Missing ⚠️
diskann-inmem/src/buffer.rs	94.71%	12 Missing ⚠️
diskann-inmem/src/layers/full.rs	96.25%	11 Missing ⚠️
diskann-inmem/src/test/epoch.rs	94.81%	10 Missing ⚠️
diskann-inmem/src/test/sequencer.rs	87.50%	6 Missing ⚠️
diskann-inmem/src/layers/mod.rs	50.00%	5 Missing ⚠️
diskann-inmem/src/freelist.rs	97.96%	4 Missing ⚠️
diskann-benchmark-runner/src/files.rs	0.00%	3 Missing ⚠️
... and 1 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1206      +/-   ##
==========================================
+ Coverage   89.64%   89.82%   +0.17%     
==========================================
  Files         503      518      +15     
  Lines       95761    98963    +3202     
==========================================
+ Hits        85848    88895    +3047     
- Misses       9913    10068     +155     

Flag	Coverage Δ
miri	89.82% <95.31%> (+0.17%)	⬆️
unittests	89.50% <95.31%> (+0.18%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-benchmark-runner/src/app.rs	83.97% <100.00%> (ø)
diskann-benchmark-runner/src/utils/fmt.rs	98.80% <100.00%> (+0.40%)	⬆️
diskann-benchmark/src/index/mod.rs	100.00% <ø> (ø)
diskann-inmem/src/counters.rs	100.00% <100.00%> (ø)
diskann-inmem/src/lib.rs	100.00% <100.00%> (ø)
diskann-inmem/src/neighbors.rs	100.00% <100.00%> (ø)
diskann-inmem/src/num.rs	100.00% <100.00%> (ø)
diskann-inmem/src/tag.rs	100.00% <100.00%> (ø)
diskann-vector/src/distance/implementations.rs	95.06% <100.00%> (ø)
diskann-inmem/src/sharded.rs	99.46% <99.46%> (ø)
... and 10 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[magdalendobson]

What is FROZEN for? Why would we want to use it?

[hildebrandmw]

I used FROZEN for the start points. Basically, it's a mechanism to mark slots a being available for reading, but prevents code from unintentionally changing them as slot retirement needs to transition from AVAILABLE to RETIRED.

There's a bigger discussion to be had about how to safely change start points and all that for a truly dynamic index that I kind of punted on for this PR.

[magdalendobson]

What would go wrong if I only ever used/updated the inline tags?

[hildebrandmw]

Used: nothing. These tags are meant to be authoritative, and so should always be correct.

Updated: lots (unless the invasive tags for the primary buffer were removed). Currently, the invasive tags are the only ones checked on the read path, so failure to update them properly would be pretty disastrous.

[magdalendobson]

I think my broader question is what do we gain from having the invasive tags? Do they make the read path faster or something?

[hildebrandmw]

The thought was basically this: for some quantization types (e.g. spherical), the overall vector length does not end up being a multiple of the cache line size anyways. Since we pad to cache line boundaries, we can basically get the invasive tag "for free", which by collocating it with the data means we don't necessarily need to generate extra bandwidth when checking the tags.

Unfortunately, this doesn't always work out for full-precision data, which often ends up being a multiple of the cache line size. However, moving to only using the tags out of the authoritative source does mean that we need to introduce an extra prefetch loop somewhere, either prefetching the tags during filtering higher in the expand_beam implementation, or emitting prefetches to the tags inline with the data access like I am currently doing.

I think the structure of Store needs to be a bit more open-ended anyways to support quantization/filtering metadata, so we might be able to make invasive tags more optional anyways.

[magdalendobson]

Does 256 guard slots mean that with default capacity we're effectively limited to 256 threads? If so this is a pretty subtle limitation and might trip people up if they are trying out DiskANN on a large machine. Would there be a downside to making the default capacity slightly larger, maybe 512?

[hildebrandmw]

No problem bumping up to 512. It does make epoch scans a little more expensive, but those should be relatively rare. I meant to expose this as part of the top level provider configuration and ... kind of forgot.

[magdalendobson]

I feel mildly uncomfortable with having an effective limit on the number of threads buried deep inside the code, maybe we should bring it up to the top level configuration after all

[hildebrandmw]

I agree. Fortunately, this is just a matter of plumbing that I hadn't bothered to do yet. Thanks for bringing it up!

[magdalendobson]

Typo :)

[hildebrandmw]

Wouldn't be a Mark PR without rampant typos in the comments.

[magdalendobson]

How do we feel about this given that 10 billion size indices may be a possibility? Is this something we want to change in the future?

[hildebrandmw]

Do we really see someone using an in-memory index for a 10-billion sized dataset? This would require TBs of memory. I'm not necessarily opposed to making this change, but I feel like at that scale, you would want a different data structure that could swap out to disk more efficiently.

[magdalendobson]

For something SIFT-ish or a little bit bigger, for a 4B sized index you need 1.5-2 TB main memory. I took a quick look at the largest VMs out there you can rent and you can definitely get up to 3 TB. I don't think it's fully unreasonable to say we just don't think this is something you should be doing with this inmem provider, though.

[magdalendobson]

We mean epochs here, not IDs, right?

Also, what is so bad about wrapping around? Couldn't we just have a special case in try_advance for handling u64::MAX?

[hildebrandmw]

I'd have to think though this - this might be a remnant from before when the actual epochs mattered. I do think there is a nasty corner case in the current implementation with the rotating queue strategy: wrapping from u64::MAX to 1 (needed to differentiate between unclaimed and claimed slots) would skip one queue in the rotation and break things upstream.

If we determine that we can support wrapping, we could use a 32-bit epoch with still 64-bit slots in the registry, which would allow us to differentiate between "taken" vs "not-taken" and epoch 0.

That said, I'm not entirely convinced this is really needed since this would need to be running for a long-long time before overflowing.