[multi vector] Add tiled MinMax MaxSim kernels

[juchen-ms (partychen)]

Summary

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

Query preparation packs the query vectors and their compensation metadata once. Each computation borrows canonical MinMax4 documents, decodes document tiles into a shared byte layout, and computes one MaxSim distance per query vector.

Data layout

Query storage

Queries use BlockTransposed<u8, MR, PACK> storage. MR is the number of query rows in a panel, and PACK is the number of consecutive bytes from each query consumed by one integer dot-product step. The packed views and panels carry PACK explicitly.

EvenOdd64Layout defines the dimension order shared by query preparation and document decoding. Within each block of 64 dimensions, even dimensions precede odd dimensions:

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

The original dimension D is rounded up to a padded dimension K, a multiple of 64. Dimension padding and unused query rows contain zero. Zero-dimensional inputs are handled separately. MinMax compensation uses the original dimension D.

PACK controls physical query grouping independently of the logical dimension permutation. The contraction length counts padded scalar dimensions, with K / PACK dot-product steps.

Document decoding

The decoder consumes four-bit codes without canonical headers or compensation metadata. Each complete block contains 32 packed bytes and produces 64 byte-valued codes: the 32 low nibbles followed by the 32 high nibbles.

• V3: A 32-byte SIMD load followed by vector masks and shifts produces the two output halves.
• V4: Uses the AVX2-capable decoder and VNNI for integer dot products.
• Neon: Two 16-byte SIMD loads produce the same 64-byte output layout.
• Scalar: Implements the same layout with portable scalar operations.

Partial blocks are zero-filled: V4 uses a masked load, and the other backends copy into a bounded local buffer. When D is odd, the unused high nibble of the final packed byte is cleared. The decoder overwrites every output byte, including padding.

Components and execution

Component	Responsibility
Prepared	Owns packed query values, query compensation, and the selected architecture capability
layout.rs	Defines the permutation, padding, and controlled packed-query storage/view construction
decode.rs	Expands packed code bytes into the shared layout using the selected architecture
reader.rs	Adapts canonical MinMax4 rows, preserves metadata order, and invokes the decoder
Driver	Validates shapes, manages scratch and cache tiling, and coordinates traversal
PanelKernel / MicroKernel	Perform integer contraction, MinMax compensation, and MaxSim reduction

The driver uses B-first traversal so decoded document data is reused across query tiles and panels:

Visit B tile
  Read canonical rows and decode values into scratch
  Visit A cache subviews
    Visit A panels
      Visit decoded B panels and remainder
        Integer contraction -> compensation -> MaxSim reduction

Traversal uses the library's packed/unpacked views, panel visitors, remainder handling, byte-based cache-blocking parameters, and bounds-tracked pointers. Cache working-set estimates include query and document compensation metadata.

The driver receives a PackedQueryView and MinMax4Rows, which retain their layout and representation contracts. Integer contraction consumes byte panels and the padded dimension; compensation separately uses metadata and the original dimension. Register representations remain opaque behind ExtraWide.

Two per-call scratch buffers hold decoded B values and B metadata for one tile and are reused across the computation. Persistent query state remains read-only. Documents are not modified, and no full intermediate distance matrix is materialized.

Architecture configurations

ISA	MR x NR	PACK	Integer computation
Scalar	8 x 6	4	Portable u32 accumulation
x86-64 V3	16 x 6	4	AVX2 maddubs and madd
x86-64 V4	16 x 8	8	AVX-512 VNNI dpbusd, with adjacent accumulator lanes folded during reduction
ARM64 Neon + dotprod	8 x 8	4	UDOT

Architecture-specific operations execute within their target-feature context. Logical query-row counts allow a backend to omit an unused upper half of a query panel.

API

let kernel = build_minmax_max_sim(MaxSimIsa::Auto, query_minmax8, BoxErase)?;
let mut scores = vec![0.0; kernel.nrows()];
kernel.compute_max_sim(doc_minmax4, &mut scores)?;

The factory supports automatic architecture selection and explicit Scalar, V3, V4, and Neon requests. Unavailable requested architectures return NotSupported. MaxSimIsa::Reference is not supported by this factory.

Behavior

• Produces one MaxSim distance per query vector and overwrites scores on each call.
• Rejects mismatched output lengths or document dimensions before modifying scores.
• Returns f32::MAX for empty documents and zero for nonempty zero-dimensional documents.
• Ignores NaN pairwise distances while preserving the best valid score.
• Keeps document scratch local to each call so prepared query state can be shared.
• Adds no dependencies; public entry points are additive.

Test coverage

The PR includes focused tests for:

• Explicit-PACK views and panels, padded columns, group spans, subviews, and tails.
• Query permutation, controlled row writes, and row/dimension padding.
• Scalar and SIMD decoding, unaligned input, nonzero unused nibbles, output sentinels, and invalid buffer lengths.
• Canonical row/metadata alignment and decoded scratch extents.
• Integer contraction against independently constructed inputs and expected results.
• Register accumulation, compensation, cache-tile variations, and single/multiple query panels.
• Factory equivalence, shape errors, empty inputs, repeated computation, prepared-query ownership, and NaN handling.

[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 98.33241% with 30 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.69%. Comparing base (fcf9053) to head (aff6533).

Files with missing lines	Patch %	Lines
...ion/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs	98.27%	15 Missing ⚠️
diskann-wide/src/arch/x86_64/v4/conversion.rs	0.00%	6 Missing ⚠️
.../matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs	97.11%	3 Missing ⚠️
diskann-quantization/src/minmax/multi/factory.rs	99.09%	3 Missing ⚠️
diskann-wide/src/arch/x86_64/v3/conversion.rs	50.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1394      +/-   ##
==========================================
+ Coverage   90.58%   90.69%   +0.11%     
==========================================
  Files         568      574       +6     
  Lines      112680   114387    +1707     
==========================================
+ Hits       102069   103742    +1673     
- Misses      10611    10645      +34     

Flag	Coverage Δ
miri	90.69% <98.33%> (+0.11%)	⬆️
unittests	90.34% <90.68%> (-0.01%)	⬇️

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

Files with missing lines	Coverage Δ
...n-quantization/src/matrix_kernels/blocks/packed.rs	100.00% <100.00%> (+0.34%)	⬆️
...quantization/src/matrix_kernels/blocks/unpacked.rs	100.00% <ø> (ø)
diskann-quantization/src/matrix_kernels/bounds.rs	100.00% <ø> (ø)
.../matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs	100.00% <100.00%> (ø)
.../matrix_kernels/maxsim/minmax8_x_minmax4/reader.rs	100.00% <100.00%> (ø)
diskann-quantization/src/minmax/multi/kernel.rs	100.00% <100.00%> (ø)
...-quantization/src/multi_vector/block_transposed.rs	98.60% <100.00%> (+0.10%)	⬆️
diskann-wide/src/emulated.rs	98.31% <ø> (ø)
.../matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs	97.11% <97.11%> (ø)
diskann-quantization/src/minmax/multi/factory.rs	99.09% <99.09%> (ø)
... and 3 more

... and 5 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)]

This is playing pretty fast and loose with the expects. Generally, panicking upwraps should be reserved for things that can only happen as logical errors in the internal implementation.

There are a few somewhat dangerous assumptions being made here:

• doc.repr().ncols() == 0: Not guaranteed because MinMaxMeta does not guarantee this in the constructor. In particular, the computation with intrinsic_dim could overflow. This invariant should be built into MinMaxMeta directly.
• Similarly, converting MatRef<'_, MinMaxMeta<_>> to a slice is something that MatRef should handle automatically.

[Mark Hildebrand (hildebrandmw)]

By popular demand, we know have a way (see #1437) to support odd-ball instructions like this. Please add this to the V3 and V4 backends with a miri fallback. Then we won't constantly have to work around it.

[Mark Hildebrand (hildebrandmw)]

This feels like a break-down in abstractions. It layers a permutation pattern on top of a data structure with a permutation pattern, making it harder to reason about the overall layout.

I think we need a better way of spinning up customizations like this. For example, kind of like PermutationStrategy but generalized. Mat was supposed to be able to do this, but the approach taken by that turned out to be pretty compromised. We need a central way of working with interleavings like this.

[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.

[Mark Hildebrand (hildebrandmw)]

Micro-kernels should be doing as regular and as minimal work as possible. Why is there an Option in the query type?

[Mark Hildebrand (hildebrandmw)]

Please add missing diskann-wide abstractions instead of working around them. It's more durable and keeps code here much less noisy.

For example, this one could use a SIMDReinterpret<u64x8, i8x16>.

[Mark Hildebrand (hildebrandmw)]

Surely there is a better way than a 10 argument macro...

[Mark Hildebrand (hildebrandmw)]

This many panics in a constructor is a sign that (A) results should be used instead (B) we're missing an abstraction somewhere else or (C) the internal invariants are too complicated.

In this situation, I think it's a heavy dose of B.

The a and b types should be able to report their number of bytes without requiring checked arithmetic, there is likely a way to package a with its metadata earlier, etc.