[multi vector] Add tiled MinMax MaxSim kernels

[juchen-ms (partychen)]

• 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?

No new dependencies are added. Existing API signatures and valid-input behavior are preserved. Invalid MinMaxMeta dimensions or byte extents are now rejected during construction; the canonical data format is unchanged.

Reference Issues/PRs

See also #1368 for the existing matrix-kernel decomposition and #1479 for the upstream diskann-wide instruction wrappers used by this implementation.

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

Adds prepared, tiled MaxSim kernels for MinMax8 queries against canonical MinMax4 documents, with Scalar, x86-64 V3/V4, and ARM64 Neon implementations.

Query values and compensation metadata are prepared once. Each computation borrows a document's quantized vectors and overwrites one compensated MaxSim distance per query vector, without materializing a full pairwise score matrix.

API

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

// query: Mat<MinMaxMeta<8>>
// document: Mat<MinMaxMeta<4>>
let kernel = build_minmax_max_sim(MaxSimIsa::Auto, query.as_view(), BoxErase)?;
let mut scores = vec![0.0; kernel.nrows()];
kernel.compute_max_sim(document.as_view(), &mut scores)?;

The factory supports automatic ISA selection and explicit Scalar, V3, V4, and Neon requests. Unavailable architectures, unsupported prepared extents, and dimensions exceeding the integer-accumulator limit of 1,122,867 return NotSupported. MaxSimIsa::Reference is not supported by this factory.

Layout and execution

The implementation uses the existing Drive, PanelKernel, and MicroKernel machinery, shared packed/unpacked views, and Slice/Elements interfaces.

• layout.rs defines the shared MinMax dimension ordering and padding contract.
• a.rs provides borrowed packed-query views with synchronized compensation metadata.
• b.rs handles canonical document-row access, scratch preparation, nibble unpacking, and prepared B views.
• Owned query storage remains in minmax/multi/query.rs, above the matrix-kernel layer.

Within each 64-dimensional block, even coordinates precede odd coordinates:

[d0, d2, ..., d62, d1, d3, ..., d63]

A uses BlockTransposed<u8, MR, PACK> storage. B is unpacked into contiguous byte-valued rows with the same dimension order. Missing dimensions and unused query rows are zero-padded; compensation uses the original dimension.

B-first traversal prepares each document tile once and reuses it across all A panels:

B tile preparation
  → A panels
    → B panels and remainder
      → integer contraction
      → MinMax compensation and reduction

The driver obtains prepared working-set sizes from AView::panel_bytes() and Scratch::row_bytes(). These include padding and compensation metadata; B's size reflects the unpacked representation, not its canonical packed row size.

Two per-call scratch buffers hold B values and metadata and are reused between tiles. Prepared query state remains read-only, and documents are not modified.

Architecture configurations

Backend	MR	NR	PACK	Integer computation
Scalar	8	5	4	Portable u32 accumulation
x86-64 V3	16	5	4	AVX2 maddubs followed by madd
x86-64 V4	16	8	8	AVX-512 VNNI with adjacent accumulator-lane folding
ARM64 Neon + dotprod	8	8	4	Unsigned dot products

Supported panel geometries and remainder instances are registered per backend. SIMD query representations are fixed-size register arrays: loading and dot accumulation always process the complete prepared patch, without Option or active-row-dependent half-panel selection.

The nibble-unpacking interface accepts arbitrary-length slices independently of SIMD register width. Backend implementations handle complete chunks and bounded tails, including the unused high nibble of odd-dimensional documents.

Validation and behavior

MinMaxMeta::try_new() validates the canonical dimension, row stride, total byte extent, and allocation limits. The validated row stride is cached. Existing MinMaxMeta::new() delegates to this validation and retains a documented panicking convenience API. Benchmark fixtures use the shared fallible constructor instead of maintaining separate size checks.

The computation:

• Rejects output-length and document-dimension mismatches before modifying scores.
• Returns f32::MAX for empty documents and zero for nonempty zero-dimensional documents.
• Handles empty queries without entering the tiled kernel.
• Ignores NaN distances while preserving the best valid score.
• Keeps valid-row handling in compensation and output handling, outside integer contraction.

Directed tests cover query packing, padding, nibble unpacking, logical group bounds, metadata alignment, prepared byte footprints, cache-tile selection, scratch reuse, unsigned accumulator conversion, and equivalence with the pairwise implementation.

Benchmark integration

Extends the existing multi-vector benchmark feature with element_type: "minmax8", quantized fixtures, regression support, example jobs, and ISA-specific performance inputs.

Quantization and query preparation are outside the timed region. The benchmark includes per-call scratch work and reports ns/IP while repeatedly scoring the same document matrix.

Any other comments?

This revision retains the shared 64-dimensional MinMax prepared format. It does not introduce a generalized column-order policy or adapter, a new Patch abstraction, or changes to shared BlockTransposed semantics.

Fixed-size query processing deliberately favors a regular micro-kernel. It can increase work for partially filled query panels and is not presented as a universal performance improvement. The choice of a shared 64-dimensional format is also a workload-dependent tradeoff, not a claim of optimality on every architecture. No native ARM performance measurements are claimed.

Local validation completed for this revision

Command	Result
cargo fmt --all	Passed
cargo clippy --workspace --all-targets -- -D warnings	Passed
cargo test -p diskann-quantization --lib minmax --quiet	Passed: 51 tests
cargo test -p diskann-benchmark --features multi-vector multi_vector --quiet	Passed: 14 tests
cargo test -p diskann-benchmark --no-default-features multi_vector --quiet	Passed: 2 tests
cargo clippy -p diskann-quantization --lib --tests --target aarch64-pc-windows-msvc -- -D warnings	Passed
cargo +nightly miri test -p diskann-quantization --lib bytes_include	Passed: 2 targeted tests

The listed Miri run used RUSTFLAGS=-Ctarget-cpu=x86-64-v3 and MIRIFLAGS=-Zmiri-disable-isolation -Zmiri-strict-provenance. These are focused checks, not a full-workspace test or Miri run.

[Copilot]

🔵 Needs a closer look

The public factory has a Reference ISA compatibility issue, and V3/V4 runtime paths remain unvalidated on x86-64.

Pull request overview

Adds prepared, tiled MaxSim kernels for MinMax8 queries against MinMax4 documents with ISA-specific dispatch.

Changes:

• Adds public prepared-kernel and factory APIs.
• Implements Scalar, AVX2/V3, AVX-512/V4, and ARM64 Neon kernels.
• Adds packing, fused compensation/reduction, validation, and reuse tests.

File summaries

File	Summary	Final review note
diskann-quantization/src/multi_vector/mod.rs	Re-exports the factory.	—
diskann-quantization/src/minmax/multi/mod.rs	Registers and exports kernel modules.	—
diskann-quantization/src/minmax/multi/kernel.rs	Defines kernel interfaces.	—
diskann-quantization/src/minmax/multi/factory.rs	Provides ISA dispatch and prepared-kernel construction.	Moderate: MaxSimIsa::Reference is reported available but rejected by this factory.
diskann-quantization/src/minmax/mod.rs	Exposes MinMax APIs.	—
diskann-quantization/src/matrix_kernels/maxsim/mod.rs	Registers the MinMax matrix kernel.	—
diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs	Implements packing, tiling, compensation, SIMD kernels, and tests.	Nit: V3/V4 paths lack x86-64 runtime coverage.

Review details

Suppressed comments (2)

diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs:785

• The V3/V4 implementations here are only compile-checked in the stated ARM64 validation: the ISA tests return early when the requested x86 ISA is unavailable, so the pdep/VNNI path and its lane reduction are not exercised. Please add or run an x86-64 runtime test (or an equivalent testable emulation) covering these kernels before relying on this new code.

    micro_kernel!(V3, 16, micro_kernel, {8, 7, 6, 5, 4, 3, 2, 1});
    micro_kernel!(V4, 16, micro_kernel, {8, 7, 6, 5, 4, 3, 2, 1});

diskann-quantization/src/minmax/multi/factory.rs:171

• MaxSimIsa::Reference is documented as an always-available selector (is_available() returns true), and the existing MinMax implementation provides the reference MaxSim path, but this new public factory rejects it unconditionally. Callers that use is_available() to preflight an ISA therefore receive NotSupported for a value reported as buildable. Please add a reference adapter (or make support for this factory explicit in the ISA API/documentation).

        MaxSimIsa::Reference => Err(NotSupported {
            isa,
            reason: "reference kernel unavailable",
        }),

• Files reviewed: 7/7 changed files
• Comments generated: 0
• Review effort level: Lite

---

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

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 99.19104% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.09%. Comparing base (3cda29e) to head (a79c823).

Files with missing lines	Patch %	Lines
diskann-benchmark/src/main.rs	97.93%	3 Missing ⚠️
diskann-benchmark/src/multi_vector/kernels.rs	95.38%	3 Missing ⚠️
diskann-quantization/src/minmax/multi/factory.rs	98.33%	3 Missing ⚠️
...ion/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs	99.61%	2 Missing ⚠️
diskann-benchmark/src/inputs/multi_vector.rs	90.90%	1 Missing ⚠️
diskann-quantization/src/minmax/multi/meta.rs	97.91%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1394      +/-   ##
==========================================
+ Coverage   91.98%   92.09%   +0.11%     
==========================================
  Files         586      593       +7     
  Lines      118318   119897    +1579     
==========================================
+ Hits       108831   110420    +1589     
+ Misses       9487     9477      -10     

Flag	Coverage Δ
miri	92.09% <99.19%> (+0.11%)	⬆️
unittests	92.04% <99.18%> (+0.11%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-benchmark/src/multi_vector/driver.rs	89.92% <100.00%> (+1.40%)	⬆️
diskann-benchmark/src/multi_vector/mod.rs	100.00% <100.00%> (ø)
diskann-benchmark/src/utils/mod.rs	95.71% <ø> (ø)
...n/src/matrix_kernels/maxsim/minmax8_x_minmax4/a.rs	100.00% <100.00%> (ø)
...n/src/matrix_kernels/maxsim/minmax8_x_minmax4/b.rs	100.00% <100.00%> (ø)
.../matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs	100.00% <100.00%> (ø)
diskann-quantization/src/minmax/multi/kernel.rs	100.00% <100.00%> (ø)
diskann-quantization/src/minmax/multi/query.rs	100.00% <100.00%> (ø)
diskann-benchmark/src/inputs/multi_vector.rs	64.78% <90.90%> (+10.69%)	⬆️
diskann-quantization/src/minmax/multi/meta.rs	99.70% <97.91%> (-0.30%)	⬇️
... and 4 more

... and 9 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[Mark Hildebrand (hildebrandmw)]

There's a lot going on here. It does add a minmax8x4 kernel, but basically builds a second matrix-kernel decomposition architecture rather than extending the existing abstractions introduced in 1368.

The point of the machinery introduced there was to use enable better higher-level experimentation of cache tilings and ordering without fully relying on Miri or checked accesses for correctness debugging. All of these pieces could then be tested, reviewed, and reused independently.

Here, packing and traversal are all reimplemented manually with inconsistent bounds tracking applied. This is hard to test, validate, or reuse.

My suggestion would be to follow the design philosophy of 1368 and build the traversal mechanism out of lower-level parts (which we can then reuse for other quantization kernels). In addition, matrix kernels should not own any non-scratch state because that also makes them much more difficult to reuse.

I also have some other design concerns:

• The microkernel is currently doing a lot, and I feel that there has to be a more efficient unpacking than doing this much bit-twiddling in general purpose registers. One design worth considering incrementally unpacking B and reusing that for the whole A traversal (flipping the cache tiling order - which again is easier to do via paneled views). This can perhaps be coupled with a different permutation strategy of the A-side query for more efficiency.
• ExtraWide should as much as possible keep the associated types opaque to the caller (like ExtraWide in the f32 kernel). This makes it significantly easier to express architecture specific optimizations of coarser grained kernels without the profoundly heaver SIMDVector and such constraints.
• How many of the Miri exceptions are really needed? Miri can emulate a good number of Neon and AVX2 intrinsics.
• Directed testing of the various steps is pretty sparse. Especially given its ad-hoc mix of checked, unchecked, and Bounds based indexing. Using Bounds based checking in a more disciplined way at least provides higher confidence of integration-based tests.

Have you considered making the A-side packing even more aggressive, splitting into groups of even and odd indices? The idea there is that it would naturally fit the order of unpacked nibbles much better and avoid a lot of the interleaving logic. Basically, a lot the micro-kernels are concerned about restoring a dimension ordering that we anyways control.

[juchen-ms (partychen)]

There's a lot going on here. It does add a minmax8x4 kernel, but basically builds a second matrix-kernel decomposition architecture rather than extending the existing abstractions introduced in 1368.

The point of the machinery introduced there was to use enable better higher-level experimentation of cache tilings and ordering without fully relying on Miri or checked accesses for correctness debugging. All of these pieces could then be tested, reviewed, and reused independently.

Here, packing and traversal are all reimplemented manually with inconsistent bounds tracking applied. This is hard to test, validate, or reuse.

My suggestion would be to follow the design philosophy of 1368 and build the traversal mechanism out of lower-level parts (which we can then reuse for other quantization kernels). In addition, matrix kernels should not own any non-scratch state because that also makes them much more difficult to reuse.

I also have some other design concerns:

• The microkernel is currently doing a lot, and I feel that there has to be a more efficient unpacking than doing this much bit-twiddling in general purpose registers. One design worth considering incrementally unpacking B and reusing that for the whole A traversal (flipping the cache tiling order - which again is easier to do via paneled views). This can perhaps be coupled with a different permutation strategy of the A-side query for more efficiency.
• ExtraWide should as much as possible keep the associated types opaque to the caller (like ExtraWide in the f32 kernel). This makes it significantly easier to express architecture specific optimizations of coarser grained kernels without the profoundly heaver SIMDVector and such constraints.
• How many of the Miri exceptions are really needed? Miri can emulate a good number of Neon and AVX2 intrinsics.
• Directed testing of the various steps is pretty sparse. Especially given its ad-hoc mix of checked, unchecked, and Bounds based indexing. Using Bounds based checking in a more disciplined way at least provides higher confidence of integration-based tests.

Have you considered making the A-side packing even more aggressive, splitting into groups of even and odd indices? The idea there is that it would naturally fit the order of unpacked nibbles much better and avoid a lot of the interleaving logic. Basically, a lot the micro-kernels are concerned about restoring a dimension ordering that we anyways control.

Mark Hildebrand (@hildebrandmw) Thanks for the detailed review. I reworked the implementation around the abstractions from #1368 in ed75906e:

• There is now a single generic Driver -> PanelKernel -> MicroKernel traversal for Scalar, V3, V4, and Neon; the separate V4/direct traversal was removed.
• Grouped<N> owns both query packing and canonical MinMax4 group expansion. Scalar/V3/Neon use Grouped<4> with even/odd dimensions, while V4 uses contiguous Grouped<8> values.
• B remains in canonical packed MinMax4 form. The same micro-kernel pipeline expands only the current contraction group, so the matrix kernel has no persistent state or document scratch allocation.
• ExtraWide keeps its register types opaque; architecture backends only provide load/splat/dot/reduce behavior.
• The Miri exceptions are limited to operations Miri cannot execute natively: BMI2 expansion, the emulated V4 broadcast, and adjacent-lane reduction. The native V4 path and the Miri-emulated V4 test both pass.

I also added/ran directed packing, tail, bounds, driver/register, factory-equivalence, repeated-computation, empty-input, and NaN-compensation coverage. Native x86-64 V3/V4 tests, ARM64 cross-Clippy, and the V4 Miri test pass.

For the 1,000-document 16x16 workload on a Xeon Platinum 8370C, the latest release measurements are:

Dim	Existing	V3	V4
250	11.47 us/doc	1.65 us/doc	0.52 us/doc
256	4.09 us/doc	1.61 us/doc	0.49 us/doc

The 250-dimensional Existing path has a scalar remainder, so I am keeping the 250- and 256-dimensional results explicitly separate.

[Mark Hildebrand (hildebrandmw)]

I got through packed.rs. I think there's a little we can do to make the group access a bit semantically tighter. I'm still working my way through the rest. Is there a concern about using the same blocking strategy across all architectures?

[juchen-ms (partychen)]

I got through packed.rs. I think there's a little we can do to make the group access a bit semantically tighter. I'm still working my way through the rest. Is there a concern about using the same blocking strategy across all architectures?

Thanks! I addressed the packed.rs comments; replies are inline.

Is there a concern about using the same blocking strategy across all architectures?

For V3, yes. I measured register and cache blocking on a Xeon Platinum 8370C (AVX-512) with release builds, across the multi-vector shapes from 16x16x256 to 64x1250x512. I didn't use #[test] timings: Bound tracking is enabled under cfg(test), and it distorts the results.

• V3: 16x8 needs 16 accumulators, all 16 ymm registers, so it spills. d9d4379 switches V3 to 16x6, like the f32 V3 kernel. It is 4-9% faster for 16/64 queries at 250/256 dimensions. 16x4 was about the same as 16x6; 8x8 and 8x12 didn't help.
• V4: 16x8 is still the best overall. 16x12 is 3-16% slower on five of six shapes and 4% faster on the largest. 32x4 and 32x6 are about 1.7x slower at 16 queries, where half of each panel is padding. V4 stays at 16x8.
• Cache blocking: the shared byte-based parameters look fine for MinMax. The default L1 budget is within 2% of the best B-tile size on five of six shapes; not tiling B is much slower, up to 1.85x on V4. a_panels_in_l2 has no measurable effect: the B-first driver visits all of A for each decoded B tile, so splitting A into L2-sized subviews doesn't change the access order.

[Mark Hildebrand (hildebrandmw)]

I have largely structural comments for this round. There is a lot of cross-talk between components. For example, EvenOddLayout64 duplicates dimensions that should already exist in a large number of other places, is AVX2 specific, and doesn't layer well with views. Lots of panic locations. The Decoder trait is hard-coupled with the EvenOdd64Layout, and ExtraWide is doing the entire micro-kernel, etc.

This feels like it's headed in the right direct, but still needs tightening of its underlying abstractions (and maybe bringing in new ones when current ones aren't doing the trick).

One big request I have is whenever there is some missing functionality in diskann-wide, please add it (or ask me to) so we can improve that library for everyone.

[Mark Hildebrand (hildebrandmw)]

Is 64 a universally good number? It means we can't use AVX-512 to do the on-the-fly unpacking and makes Arm work a little harder.

[juchen-ms (partychen)]

Agreed—64 is not universally optimal. It originally matched a 32-byte AVX2 load of packed 4-bit values, producing two 32-byte channels for the even/odd layout.

I compared 32-, 64-, and 128-dimensional blocks on a Xeon Platinum 8370C, keeping V4’s MR16/NR8/PACK8 unchanged. The 128-dimensional variant used native 64-byte AVX-512 decoding, rather than just increasing the padding. Across 22 multi-vector shapes, with three repeated runs, 32-dimensional blocks reduced geometric-mean latency by about 1.7% relative to 64, while 128-dimensional blocks increased it by about 5.8%. Smaller blocks helped some tail dimensions; larger blocks introduced more padded computation.

For now, I’d keep 64 as the shared prepared-data format: the measured benefit does not yet justify introducing architecture-specific layouts. This is a tradeoff supported by this x86 workload, not a claim that 64 is best everywhere. I haven’t benchmarked this on native ARM hardware, so your concern about Neon remains open.