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.

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.

Work in progress

• A/B benchmark against main to confirm the restructure hasn't cost any perf.

[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 90.92123% with 68 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.50%. Comparing base (3218478) to head (c713e51).

Files with missing lines	Patch %	Lines
...ization/src/multi_vector/distance/kernels/float.rs	87.31%	35 Missing ⚠️
...ntization/src/multi_vector/distance/kernels/f16.rs	74.35%	20 Missing ⚠️
...ntization/src/multi_vector/distance/kernels/mod.rs	92.39%	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.50%   -0.06%     
==========================================
  Files         522      521       -1     
  Lines       99541    99641     +100     
==========================================
+ Hits        91139    91176      +37     
- Misses       8402     8465      +63     

Flag	Coverage Δ
miri	91.50% <90.92%> (-0.06%)	⬇️
unittests	91.18% <90.92%> (-0.06%)	⬇️

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.00% <92.39%> (-7.00%)	⬇️
...ntization/src/multi_vector/distance/kernels/f16.rs	70.73% <74.35%> (-29.27%)	⬇️
...ization/src/multi_vector/distance/kernels/float.rs	87.31% <87.31%> (ø)

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