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.51% -0.05%
==========================================
Files 522 521 -1
Lines 99541 99614 +73
==========================================
+ Hits 91139 91158 +19
- Misses 8402 8456 +54
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.
|
|
||
| /// Fold `acc`'s columns into the running per-A-row maxima in `state`, one A-panel wide. | ||
| #[inline(always)] | ||
| pub(crate) fn fold_columns(arch: V3, acc: &[f32], state: &mut [f32]) { |
There was a problem hiding this comment.
This spills pretty heavily to the stack (all 8 vector registers get spilled). Generally, trying to rely on LLVM to correctly unroll loops like this where things are in arrays is a recipe for stack spilling. I had better luck just manually unrolling the implementation (shaving 256 bytes of stack space).
Also a nit: fold_columns is a pretty generic name for a rather specific operation.
| } | ||
| } | ||
|
|
||
| pub(super) struct QueryPanels<'a, T, const AR: usize>(ChunksExact<'a, T>); |
There was a problem hiding this comment.
So, one thing to consider that came up while I was inspecting the generated code. These iterators (while better than raw pointer soup) kind of introduce a death by a thousand cuts when it comes to code generation.
ChunksExacthas a stack footprint of 40 bytes. But forQueryPanels, we know the remainder is not used. We could easily drop this to 24 bytes (and maybe even 16 bytes ifkwere supplied externally).DocPanelshas it worse since it has aChunksExactand akwhen theTailIteratorprotocol makes it possible to derive the remainder naturally rather than eagerly materializing it.- We could shave 8 bytes off
Cursorby not making it resettable and instead creating a newCursoron demand. Tile walkers need not be resettable if we can cheaply remake them. Slotshas a slice when its size is statically known (shaving another 8 bytes)Slotssuffers the same size bloat caused byChunksExact.
Now, LLVM may be able to elide some of these fields. But all of this adds up to >700 bytes of stack space for the V3 kernel (or around 450 bytes after fixing fold_columns) with lots of messing around with the stack.
This isn't necessarily a blocker for this PR, but it is something to consider for the overall flow. The combination of extra stack space + panicking code-paths can affect overall code generation. While the code for the innermost micro-kernel looks good, it's the surrounding bit that worries me.
Usually stuff like this doesn't matter so much, but for a matrix/matrix style kernel, these optimizations can add up.
|
|
||
| /// One block-transposed block: `AR` rows × `k` columns, column-major within the block. | ||
| #[derive(Clone, Copy)] | ||
| pub(super) struct QueryPanel<'a, T, const AR: usize>(&'a [T]); |
There was a problem hiding this comment.
Can we use more descriptive names? A QueryPanel is how this happens to be used, but not what it is. Maybe BlockTransposedPanel instead?
| } | ||
| } | ||
|
|
||
| /// A run of whole blocks — block-transposed storage is padded to `AR`, hence [`NoTail`]. |
There was a problem hiding this comment.
Nit: the NoTail comment probably belongs elsewhere?
| /// Panics if `docs` has zero-length rows. | ||
| pub(super) fn new(docs: MatRef<'a, Standard<T>>, b_panels: usize) -> Self { | ||
| let k = docs.vector_dim(); | ||
| assert!(k > 0, "DocWalk requires a non-empty contraction"); |
There was a problem hiding this comment.
Suggestion: Rust has the wonderful NonZeroUsize type. Use it internally. This removes needing to constantly recheck it - and Rust/LLVM know that it cannot be zero, so divides in the future won't emit code to handle the divide-by-zero case.
| @@ -0,0 +1,104 @@ | |||
| // Copyright (c) Microsoft Corporation. All rights reserved. | |||
| // Licensed under the MIT license. | |||
There was a problem hiding this comment.
Nit: why is the header style in the distances module different from the rest of the repo?
| //! Rows are query rows and one column is one doc — note that a doc arrives as a *row* of | ||
| //! the input and lands as a *column* here. Column-major over the whole strip: column `c` | ||
| //! occupies `[c * AR, (c + 1) * AR)`, and slot `p` covers columns `p * BR ..`, so a drain | ||
| //! can address a run of columns without knowing which slot produced them. |
There was a problem hiding this comment.
A picture would make this infinitely clearer. And why not call it row major with the axes flipped?
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT license. | ||
|
|
||
| //! The accumulator a fill writes and a drain reads. |
There was a problem hiding this comment.
This is only for one particular Fill and Drain combination though, right? The documentation here makes it sound like a global property.
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 five 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.
Performance
Against
mainon 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:
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.