Skip to content

[multi-vector] Add MinMax MaxSim kernel factory - #1508

Open
juchen-ms (partychen) wants to merge 1 commit into
microsoft:mainfrom
partychen:juchen-microsoft-minmax-factory
Open

juchen-ms (partychen) wants to merge 1 commit into
microsoft:mainfrom
partychen:juchen-microsoft-minmax-factory

Conversation

@partychen

@partychen juchen-ms (partychen) commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor
  • Does this PR have a descriptive title that could go in our release notes?
  • Does this PR add any new dependencies?
  • Does this PR modify any existing APIs?
  • Is the change to the API backwards compatible?
  • Should this result in any changes to our documentation, either updating existing docs or adding new ones?

Reference Issues/PRs

Follow-up to #1494, which introduced the standalone MinMax matrix kernel. This PR is based on its merged commit and adds the factory integration without repeating the kernel's original commits.

What does this implement/fix? Briefly explain your changes.

Expose the MinMax matrix kernel through an owned-query MaxSim factory:

  • Add build_minmax_max_sim, accepting QueryMatRef<MinMaxMeta<A_BITS>>. Query and document bit widths independently support 1 through 8. Prepare the query once and reuse the kernel across borrowed documents and output buffers.
  • Support automatic backend selection and explicit Scalar, x86-64 V3/V4, and Neon requests, with layouts chosen for each document bit width. Unavailable ISAs return an error rather than silently falling back.
  • Keep ReferenceKernel as an independent scalar correctness baseline. It does not delegate to the SIMD-capable pairwise API; existing pairwise registrations and MaxSimIsa semantics remain unchanged.
  • Return one negated maximum inner product per query row. Preserve empty-input and nonfinite-score behavior, and validate document dimensions and output lengths before modifying scores.
  • Reuse BuildMaxSimError, adding Overflow for storage-size failures. Implement NewCloned for MinMaxMeta so the reference kernel uses the existing to_owned() mechanism.
  • Add focused tests for all 64 bit-width combinations and panel tails, backend/layout selection, query ownership, input validation, empty inputs, nonfinite scores, storage overflow, and cloned-storage independence. Remove the kernel's temporary dead-code suppression now that it is integrated.

No new dependencies. Existing factory signatures remain unchanged; the new entry points and the additional variant on the already non-exhaustive BuildMaxSimError are additive. Rustdoc includes the new API's contract and a usage example.

Usage example: MinMax8 query against MinMax4 documents

The inputs below are already quantized and must have matching intrinsic dimensions. Build once per query and reuse both the kernel and score buffer across documents:

use diskann_quantization::{
    minmax::{MinMaxMeta, build_minmax_max_sim},
    multi_vector::{BoxErase, Mat, MaxSimIsa, QueryMatRef},
};

fn score_documents(
    query: QueryMatRef<'_, MinMaxMeta<8>>,
    documents: &[Mat<MinMaxMeta<4>>],
) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
    let kernel = build_minmax_max_sim::<8, 4, _>(
        MaxSimIsa::Auto,
        query,
        BoxErase,
    )?;
    let mut scores = vec![0.0; kernel.nrows()];
    let mut results = Vec::with_capacity(documents.len());

    for document in documents {
        kernel.compute_max_sim(document.as_view(), &mut scores)?;
        results.push(scores.clone());
    }
    Ok(results)
}

If the query is an owned Mat<MinMaxMeta<8>>, pass QueryMatRef::from(query.as_view()). Auto selects the available backend; BoxErase packages it behind a common kernel interface. Each result contains one negated maximum inner product per query row (lower is better). Each call overwrites scores; the example clones the buffer only to retain each document's result.

Any other comments?

Validated locally on Windows x86-64 at commit 9b16ba1f8aa050e5706272f6328ed5c7dc4cd691:

Command Result
cargo fmt --all Passed
cargo clippy --workspace --all-targets -- -D warnings Passed
cargo test -p diskann-quantization --lib minmax --quiet Passed: 65 tests
cargo test -p diskann-quantization --lib --config 'profile.test.debug-assertions=false' --config 'profile.test.overflow-checks=false' minmax::multi::factory --quiet Passed: 9 tests, including overflow under release-equivalent arithmetic settings
cargo test -p diskann-quantization --no-default-features --lib minmax::multi::factory --quiet Passed: 8 tests
cargo test -p diskann-quantization --doc build_minmax_max_sim --quiet Passed: 1 doctest
cargo check -p diskann-quantization --tests --target aarch64-pc-windows-msvc --quiet Passed: compile-only
git diff --check Passed before commit

Runtime validation covers only the backends available on the local host. AArch64 execution, ISA emulation, and coverage reporting were not run locally; CI results should be checked separately.

Integrate the merged MinMax matrix kernel with an owned-query factory, runtime ISA selection, an independent scalar reference, and shared construction errors. Preserve QueryMatRef inputs and add focused factory and query-cloning coverage.

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

Copilot-Session: 94d30f43-fadc-4a6b-889b-406de424a1ab
@partychen
juchen-ms (partychen) requested review from a team and a balanced review from Copilot October 23, 2026 05:05

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Unsafe storage cloning and runtime-selected SIMD paths warrant human review with cross-platform validation results, which were not verified here.

No open issues 🎉

What changed in this PR

Integrates the MinMax matrix kernel into DiskANN’s multi-vector distance API through an owned-query factory, preserving QueryMatRef inputs.

Changes:

  • Adds runtime backend selection, prepared kernels, and an independent scalar reference.
  • Introduces shared construction errors and owned MinMax query cloning.
  • Adds coverage for bit widths, backend layouts, ownership, validation, and edge-case scores.
File Description
diskann-quantization/​src/​multi_vector/​distance/​factory.rs Adds overflow to shared construction errors.
diskann-quantization/​src/​minmax/​multi/​reference.rs Adds an independent scalar scoring baseline.
diskann-quantization/​src/​minmax/​multi/​prepared.rs Wraps packed kernels with validation and distance conversion.
diskann-quantization/​src/​minmax/​multi/​mod.rs Registers and exports the new factory interfaces.
diskann-quantization/​src/​minmax/​multi/​meta.rs Implements owned query cloning and cloning tests.
diskann-quantization/​src/​minmax/​multi/​kernel.rs Defines kernel and erasure traits with shared validation.
diskann-quantization/​src/​minmax/​multi/​factory.rs Implements backend selection, construction checks, and factory tests.
diskann-quantization/​src/​minmax/​mod.rs Exports the public MinMax factory API.
diskann-quantization/​src/​matrix_kernels/​maxsim/​packed_minmax_x_unpacked_minmax.rs Exposes the accumulator limit for factory validation.
diskann-quantization/​src/​matrix_kernels/​maxsim/​mod.rs Removes the integrated kernel’s dead-code suppression.

Review effort: Balanced

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.94521% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.14%. Comparing base (8a9ff2d) to head (9b16ba1).

Files with missing lines Patch % Lines
diskann-quantization/src/minmax/multi/factory.rs 98.00% 6 Missing ⚠️
diskann-quantization/src/minmax/multi/reference.rs 92.30% 3 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1508      +/-   ##
==========================================
- Coverage   92.16%   91.14%   -1.02%     
==========================================
  Files         592      596       +4     
  Lines      123582   124019     +437     
==========================================
- Hits       113901   113039     -862     
- Misses       9681    10980    +1299     
Flag Coverage Δ
miri 91.14% <97.94%> (-1.02%) ⬇️
unittests 90.85% <97.92%> (-1.26%) ⬇️

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

Files with missing lines Coverage Δ
..._kernels/maxsim/packed_minmax_x_unpacked_minmax.rs 94.63% <100.00%> (-3.56%) ⬇️
diskann-quantization/src/minmax/multi/kernel.rs 100.00% <100.00%> (ø)
diskann-quantization/src/minmax/multi/meta.rs 100.00% <100.00%> (ø)
diskann-quantization/src/minmax/multi/prepared.rs 100.00% <100.00%> (ø)
...-quantization/src/multi_vector/distance/factory.rs 77.51% <ø> (ø)
diskann-quantization/src/minmax/multi/reference.rs 92.30% <92.30%> (ø)
diskann-quantization/src/minmax/multi/factory.rs 98.00% <98.00%> (ø)

... and 52 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@partychen juchen-ms (partychen) changed the title Add MinMax MaxSim kernel factory [multi-vector] Add MinMax MaxSim kernel factory Oct 23, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants