Redesign the multi-vector kernels to decouple layout, micro-kernel and reduction - #1333
Conversation
There was a problem hiding this comment.
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 genericdriveloop 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 Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1333 +/- ##
==========================================
- Coverage 91.55% 91.54% -0.02%
==========================================
Files 522 520 -2
Lines 99541 99704 +163
==========================================
+ Hits 91139 91270 +131
- Misses 8402 8434 +32
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Mark Hildebrand (hildebrandmw)
left a comment
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks - I took one pass over this. This is a step in the right direction, but there is one large theme that I think would really help a future review cycle, and that is documentation. More concise and self-contained documentation would go a long way toward making this easier to review and maintain.
Left hand arguments are referred to as query and A while the right is referred to as document or B, and these are conflated throughout the stack. I'd recommend sticking to just A and B for everything but the uppermost layers.
As an example, the documentation for QueryTile is
A run of whole blocks — block-transposed storage is padded to `AR`, hence [`NoTail`].
But rewriting as
A view over consecutive blocks from a [`BlockTransposedRef<T, AR, 1>`].
Its [`Paneled`] implementation yields one [`QueryPanel`] per block.
with the following on QueryPanel:
A single sub-block of a [`BlockTranpose`] containing `AR` rows
in a **column-major** layout.
ties the implementation to the logical data structure its operating over (and it's corresponding documentation), while providing enough breadcrumbs that someone reading the code here can piece things together a little more.
Thanks Mark Hildebrand (@hildebrandmw) for the guidance on the documentation, tried to fix all. |
Mark Hildebrand (hildebrandmw)
left a comment
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks Suryansh. I've been staring at this for a while trying to puzzle though how everything works and I'm still struggling.
X buckets of concerns:
- The performance loss outlined in the PR description seems almost entirely self-inflicted. This redesign introduces the
Stripfor the accumulator (which aggregates partial products like a normal matrix multiplication) anddrainsit out to the final result. But it seems like this design is entirely compatible with accumulating into the result like the old one? Or at least eagerly compressing the maximum elements when writing back to the accumulator. Doing so would considerably reduce the space needed for the accumulator and probably get back some or all of the performance. It would also resolve the semantics of "rows" and "columns" inStripas it would just become one-dimensional. - Lots of code is duplicated between the v3 and the scalar kernel. It's not the end of the world but does stand out to me.
- There is a non-locality that is hard to follow. It's better than dancing around with raw pointers everywhere, but still feels like you need to know how everything works to know anything works.
I understand that this last point is particularly vague and it could be that everything will click tomorrow when I come back fresh. And it could be that this is just an intrinsic property of writing dense kernels.
While not necessarily a blocker, this property in the code and particularly the documentation (which often focuses on how something fits into some global picture rather than describing how a struct works and the invariants it needs/maintains) it making it difficult for me to conceptualize how everything fits together.
| /// Merge each column of `acc` into the running per-A-row maxima in `state`. | ||
| /// | ||
| /// The maxima are held in registers for the whole sweep, so `state` is read once on entry | ||
| /// and written once on exit, not once per column. |
There was a problem hiding this comment.
For documentation of internal methods like this, please state what it actually does, not how it's used. For example:
Computes `state[a] = max(state[a], acc[0][a], acc[1][a], ...)`
is much clearer.
|
|
||
| /// Merge up to [`WAYS`] consecutive columns, one per chain. | ||
| #[inline(always)] | ||
| fn max_into_chains(arch: V3, chains: &mut [[f32s; REGS]; WAYS], src: &[[f32; A_PANEL]]) { |
There was a problem hiding this comment.
You could probably use as_flattened to simplify this considerably.
| /* | ||
| * Copyright (c) Microsoft Corporation. | ||
| * Licensed under the MIT license. | ||
| */ |
There was a problem hiding this comment.
This seems copied almost verbatim from v3.rs with small tweaks: different panel sizes and no FMA. Is that really the best way to structure it?
|
|
||
| /// Independent max chains in `max_into_rows`: enough to keep a multi-cycle max off its own | ||
| /// critical path, few enough that the chains stay in registers. | ||
| const WAYS: usize = 4; |
There was a problem hiding this comment.
Maybe move this closer to its use?
| //! | ||
| //! Panel geometry is fixed per ISA, not derived from a lane count. The scalar leaf is | ||
| //! deliberately narrower than `2 × LANES` would suggest, because its "lanes" are a | ||
| //! loop, not silicon. |
There was a problem hiding this comment.
This is misleading. Panel geometry is derived from a lane count (literally: pub(crate) const A_PANEL: usize = 2 * f32s::LANES). While yes this is informed formed by the ISA, I don't see why the negative is needed. Would it not be clearer to say "Panel geometry is optimized for each ISA"?
The comment about scalar is also strange. It does not use 2 x LANES.
While I may be harping on documentation, realize that as a review, I'm reading the documentation to get an understanding for a change. If the documentation leaves me more confused than before, the review gets considerably harder.
| //! loop, not silicon. | ||
| //! | ||
| //! Every leaf here **stores** its products into a slot and leaves the reduction to | ||
| //! `max_into_rows`. That reduction is not negligible work when the strip is wide and |
There was a problem hiding this comment.
Question: this PR claims that it regresses performance because of this choice. However ... we have all the pieces to fuse the max directly into the accumulator, right? Why not do that and not have a performance regression?
|
|
||
| /// Plan, allocate the strip, and drive. | ||
| /// | ||
| /// `k` is the *physical* row length both walks stride by, so it must be A's padded column |
There was a problem hiding this comment.
Why mention column padding when this kernel doesn't accept block-transposes with PACK > 1?
| /// row. A B row therefore arrives as a *row* of the input and lands as a *column* of the | ||
| /// accumulator. | ||
| /// | ||
| /// Carries `nd` because a strip's trailing columns belong to B rows past the end of the |
There was a problem hiding this comment.
"Belong to B rows past the end of the matrix". How is it belonging to rows that don't exist?
| self.0.next().map(|data| RowMajorTile::new(data, k)) | ||
| } | ||
|
|
||
| fn reset(&mut self) { |
There was a problem hiding this comment.
The refactor removes the old asymmetric multi-tile f16 coverage, while the replacement tiny-budget tests only exercise the f32 walks. This leaves the new reusable widening buffer and RowMajorWiden::reset path untested. Please add forced multi-tile f16 Scalar and V3 cases against the naive reference, including asymmetric A/B tile counts.
| // `run` seeds the max itself, so the fill value here is arbitrary. | ||
| let mut state = vec![0.0; self.prepared.padded_nrows()]; | ||
| self.arch | ||
| .run3(MaxIp, self.prepared.reborrow(), doc, &mut state); |
There was a problem hiding this comment.
This call no longer performs the old query.ncols() == doc.vector_dim() boundary check. For non-empty contractions, a mismatch reaches the opaque "B panel extent" assertion; for zero-dimensional queries or empty documents, the early returns accept the mismatch and return scores. Please validate dimensions before either early return for both f32 and f16, and report both values.
| // Compute side // | ||
| ////////////////// | ||
|
|
||
| /// One A-panel × one B-panel → one accumulator slot. |
There was a problem hiding this comment.
Because Scratch is reused, an Accumulate implementation must overwrite every live output cell represented by B; leaving stale slot data would silently corrupt the following drain. Please document this full-overwrite requirement and clarify that unused tail cells may remain untouched.
| /// | ||
| /// # Panics | ||
| /// | ||
| /// Panics if `state` is shorter than A's padded row count: the drain indexes it one |
There was a problem hiding this comment.
The # Panics section omits incompatible A/B contraction lengths. The generated walks can carry different k values, which reaches the leaf's "B panel extent" assertion. Please document this requirement here; caller-facing validation should reject it with the actual dimensions.
Rebuilds
diskann-quantization/src/multi_vector/distance/kernels/around a contract insteadof 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.rswas 806 lines in which the cachetiling, 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:
drivefunction owns the loop over blocks, written once for everybodyA 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
unsafethrough 6 of them, because the layout helpers, the tiling loop and the leaves each did their
own address arithmetic. Now
unsafeappears in two files only, the two leaves, as three blockseach, 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 anddrivetiles.rs, turning a matrix into blocks and panelsstrip.rs, the accumulator that leaves write into and drains read out ofleaves/v3.rsandleaves/scalar.rs, the two micro-kernelsfloat.rs, the f32 instantiation and its testsf16.rs, the widening walksWorth knowing going in:
driveonly 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.
Six files under
distance/appear in the diff with nothing but a licence header change. Themodule was using
//headers where the rest of the repo uses/* */. Fixed while in here.Performance
Against
mainon a shared machine, 30 alternating rounds per side, taking the best round ofeach cell. Median +1.25%, p90 +4.45%.
Eight of the nine shapes land between +0.44% and +2.09%, at or near the noise floor.
The ninth, the smallest, does not. At 8 queries x 32 docs x dim 128:
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 theseam, 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.
Switch revision, rebuild, run again into
after.json, then: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_measurementsor a looser tolerance is not a substitute.The input carries
referencerows, the same code on both sides. Their spread is your errorbar; anything smaller is unmeasured.
Review order
kernels/mod.rsfirst. It's the contract; everything else implements it.kernels/tiles.rskernels/strip.rs. Short, but mind the axes: a doc is a row of the input and a column here.kernels/leaves/. All the unsafe in the PR, ~120 lines each. Worth reading the SAFETYcomments properly.
kernels/float.rs, where it comes together, plus the test suite.kernels/f16.rs, the proof it composes: a second element type, no new leaf.factory.rs, the wiring and the only caller-facing diff.