Redesign the multi-vector kernels to decouple layout, micro-kernel and reduction

[Suryansh Gupta (suri-kumkaran)]

Rebuilds diskann-quantization/src/multi_vector/distance/kernels/ around a contract instead
of a monolith. Same results, same public API, same tests. Only the internals change.

Big thanks to Mark Hildebrand (@hildebrandmw) for the valuable design insights and constant support in shaping and refining this design.

Why

The reduction was baked into the kernel. tiled_reduce.rs was 806 lines in which the cache
tiling, the inner FMA loop and the max-reduction were one function body, written out once per
element type per instruction set. To feed it quantized vectors, or compute anything other than
MaxSim, there was no seam to cut along. You copied the file.

So the redesign cuts it into four pieces, each with an obvious owner:

• a walk turns a matrix into blocks small enough to sit in cache
• a leaf is the SIMD micro-kernel that multiplies one pair of those blocks
• a drain decides what the finished accumulator means (for MaxSim, the max down each query row)
• one drive function owns the loop over blocks, written once for everybody

A new element type is now just a new walk, a new reduction just a new drain, and neither
touches a leaf or the loop. f16 already shows this: it has no kernel code left at all, just a
walk that widens to f32 on the way in. It used to be a separate path.

The other reason is safety. Raw pointers ran through 5 of the 8 old files, and unsafe
through 6 of them, because the layout helpers, the tiling loop and the leaves each did their
own address arithmetic. Now unsafe appears in two files only, the two leaves, as five blocks
each, every one wrapping a single load, read or store.

What

Gone: tiled_reduce.rs, layouts.rs, reduce.rs, f32/{mod,scalar,v3}.rs. In their place,
roughly in dependency order:

• mod.rs, the four traits above plus the cache planner and drive
• tiles.rs, turning a matrix into blocks and panels
• strip.rs, the accumulator that leaves write into and drains read out of
• leaves/v3.rs and leaves/scalar.rs, the two micro-kernels
• float.rs, the f32 instantiation and its tests
• f16.rs, the widening walks

Worth knowing going in: drive only hands out ordinals, like "A-panel 3, B-panels 8..12",
never a stride or an address. That's what lets a drain whose panels are a different width
reuse the same loop.

One behaviour change to flag. Budgets and panel geometry are unchanged, but the B-panel count
now charges the accumulator strip against L1, which the old planner never counted. More
accurate, though it shrinks the B tile at small dims: 31 panels to 24 at dim 64, 13 to 12 at
dim 128, no change from 256 up.

Performance

Against main on a shared machine, 16 runs per side, comparing the minimum of each. Median +0.59%, p90 +2.65%.

Eight of the nine shapes land between +0.02% and +1.52%, which is inside the noise floor.
The ninth, the smallest, does not. At 8 queries x 32 docs x dim 128:

• f16 v3, 1106us to 1202us, +8.7%
• f16 v4, 1166us to 1254us, +7.6%
• f32 scalar, 667us to 705us, +5.7%
• f32 v3, 809us to 846us, +4.6%

That is the shape with the least work to amortise a fixed cost, and there is a new fixed
cost. The old kernel fused the multiply and the reduction, so the accumulator lived in
vector registers and folded into the caller's scratch on the way out; it was never
materialised. Splitting accumulate from drain means it has to exist somewhere the drain can
read it, which is one vec![0.0f32; plan.strip_len()] per call. That is the price of the
seam, paid once per call however much work follows.

Fixable, but not here: the strip is a fixed-size buffer with a known bound, so it can be
hoisted to the caller or held on the stack. Left as a follow-up to keep this PR a
restructure.

Running it

The benchmark is behind a feature flag. Run each side, then compare.

cargo run --release -p diskann-benchmark --features multi-vector -- run --input-file diskann-benchmark/perf_test_inputs/multi-vector.json --output-file before.json

Switch revision, rebuild, run again into after.json, then:

cargo run --release -p diskann-benchmark --features multi-vector -- check run --tolerances diskann-benchmark/perf_test_inputs/multi-vector-tolerance.json --input-file diskann-benchmark/perf_test_inputs/multi-vector.json --before before.json --after after.json

Reading it

One run per side is not enough. A shared machine drifts between faster and slower states and
holds each for longer than a measurement takes, so a single pair of runs will breach the
tolerance in places with no code change at all. The drift is one-sided, it can only make
things slower, so repeat both sides and take the minimum, and do enough runs that the numbers
stop moving. A bigger num_measurements or a looser tolerance is not a substitute.

The input carries reference rows, the same code on both sides. Their spread is your error
bar; anything smaller is unmeasured.

Review order

1. kernels/mod.rs first. It's the contract; everything else implements it.
2. kernels/tiles.rs
3. kernels/strip.rs. Short, but mind the axes: a doc is a row of the input and a column here.
4. kernels/leaves/. All the unsafe in the PR, ~120 lines each. Worth reading the SAFETY
   comments properly.
5. kernels/float.rs, where it comes together, plus the test suite.
6. kernels/f16.rs, the proof it composes: a second element type, no new leaf.
7. factory.rs, the wiring and the only caller-facing diff.

[Copilot]

Pull request overview

Refactors the diskann-quantization multi-vector MaxSim implementation to separate concerns (tiling/walks, micro-kernel leaves, accumulator scratch layout, and reduction/drain) behind a trait-based “driver + contracts” design, while keeping the same public API and test intent.

Changes:

• Replaces the previous monolithic tiled_reduce + layout/conversion machinery with a generic drive loop and small, composable traits (TileWalk/Paneled/Scratch/Accumulate/Drain).
• Introduces new tiling/walk primitives (tiles.rs), accumulator storage (strip.rs), and ISA-specific leaves (leaves/*) and wires them into the f32 pipeline (float.rs).
• Reworks the f16 path to widen per-tile into reusable buffers via lending walks, reusing the f32 pipeline without f16-specific leaves.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
diskann-quantization/src/multi_vector/distance/kernels/mod.rs	Defines the new kernel “contract” traits, plan, and drive loop; exports new f32/f16 entries.
diskann-quantization/src/multi_vector/distance/kernels/tiles.rs	Adds tile/panel abstractions and lending walks for block-transposed queries and row-major docs.
diskann-quantization/src/multi_vector/distance/kernels/strip.rs	Adds column-major accumulator strip partitioned into fixed-size slots for leaf writes.
diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs	Introduces per-ISA leaf module structure and reduction-chain configuration.
diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs	New AVX2+FMA f32 leaf micro-kernel and column-fold reduction.
diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs	New scalar/emulated f32 leaf micro-kernel and column-fold reduction.
diskann-quantization/src/multi_vector/distance/kernels/float.rs	New f32 MaxSim pipeline: plans, allocates strip, drives, and drains to per-row maxima + tests.
diskann-quantization/src/multi_vector/distance/kernels/f16.rs	Replaces f16 adapter with per-tile widening walks that feed the f32 pipeline.
diskann-quantization/src/multi_vector/distance/factory.rs	Switches factory dispatch from old F32Kernel/F16Entry to new MaxIp/MaxIpF16 entries.
diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs	Removed: old 5-level tiling loop implementation and its tests.
diskann-quantization/src/multi_vector/distance/kernels/layouts.rs	Removed: old layout marker + tile-level conversion traits.
diskann-quantization/src/multi_vector/distance/kernels/reduce.rs	Removed: old compile-time reduction helper trait.
diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs	Removed: old f32 kernel family entry and dispatch wrapper.
diskann-quantization/src/multi_vector/distance/kernels/f32/scalar.rs	Removed: old scalar micro-kernel implementation.
diskann-quantization/src/multi_vector/distance/kernels/f32/v3.rs	Removed: old v3 micro-kernel implementation.

---

💡 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 91.80328% with 60 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.51%. Comparing base (3218478) to head (7a2abcc).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
...ization/src/multi_vector/distance/kernels/float.rs	88.57%	28 Missing ⚠️
...ntization/src/multi_vector/distance/kernels/f16.rs	75.00%	19 Missing ⚠️
...ntization/src/multi_vector/distance/kernels/mod.rs	93.51%	7 Missing ⚠️
...ization/src/multi_vector/distance/kernels/tiles.rs	95.80%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1333      +/-   ##
==========================================
- Coverage   91.55%   91.51%   -0.05%     
==========================================
  Files         522      521       -1     
  Lines       99541    99614      +73     
==========================================
+ Hits        91139    91158      +19     
- Misses       8402     8456      +54     

Flag	Coverage Δ
miri	91.51% <91.80%> (-0.05%)	⬇️
unittests	91.19% <91.80%> (-0.05%)	⬇️

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

Files with missing lines	Coverage Δ
...-quantization/src/multi_vector/distance/factory.rs	79.86% <100.00%> (-0.41%)	⬇️
...src/multi_vector/distance/kernels/leaves/scalar.rs	100.00% <100.00%> (ø)
...ion/src/multi_vector/distance/kernels/leaves/v3.rs	100.00% <100.00%> (ø)
...ization/src/multi_vector/distance/kernels/strip.rs	100.00% <100.00%> (ø)
...ization/src/multi_vector/distance/kernels/tiles.rs	95.80% <95.80%> (ø)
...ntization/src/multi_vector/distance/kernels/mod.rs	93.96% <93.51%> (-6.04%)	⬇️
...ntization/src/multi_vector/distance/kernels/f16.rs	71.25% <75.00%> (-28.75%)	⬇️
...ization/src/multi_vector/distance/kernels/float.rs	88.57% <88.57%> (ø)

... and 44 files with indirect coverage changes

🚀 New features to boost your workflow:

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