Skip to content

Fuse MaxSim reduction at panel boundaries - #1365

Closed
juchen-ms (partychen) wants to merge 1 commit into
microsoft:users/suryangupta/multi_vector_kernel_new_designfrom
partychen:juchen-microsoft-fuse-maxsim-panel-reduction
Closed

Fuse MaxSim reduction at panel boundaries#1365
juchen-ms (partychen) wants to merge 1 commit into
microsoft:users/suryangupta/multi_vector_kernel_new_designfrom
partychen:juchen-microsoft-fuse-maxsim-panel-reduction

Conversation

@partychen

Copy link
Copy Markdown
Contributor

Context

This draft builds directly on #1333 at 4683fd8732988cbc9833186d55f1530b4c3b53b0 and targets users/suryangupta/multi_vector_kernel_new_design. It preserves that PR's separation of layout, tiling, and leaf micro-kernels while exploring a narrower reduction boundary for MaxSim.

Mark observed that materializing every raw A-panel x B-panel product into Strip and draining it later adds avoidable traffic even though the leaf already knows when a complete panel result is available. This follow-up keeps the parent architecture and refines only that handoff.

Problem

The current Scratch / Slots / Accumulate / Drain contract writes tile-sized raw products and then makes a second pass to reduce them. Besides the extra memory traffic, the contract spreads MaxSim-specific behavior across non-local plumbing, and Scalar and V3 duplicate substantial storage/reduction machinery.

Design

  • Replace the strip contract with a typed PanelOp invoked after each full A-panel x B-panel contraction.
  • Let the f32 Scalar and V3 leaves fold completed document lanes directly into operation-owned running MaxSim state.
  • Keep f16 widening its tiles and reusing the f32 path.
  • Remove kernels/strip.rs and the tile-sized intermediate allocation.
  • Keep panel geometry, tile traversal, layout types, and micro-kernel selection unchanged.

This is an architecture/performance remediation of the parent refactor, not a new kernel implementation.

Correctness

The panel operation receives the same ordinal A/B panel positions as the previous drive contract, preserving full panels, B tails, A/B tile ordinal behavior, empty dimensions/documents, and padded query rows.

The reduction retains Rust f32::max behavior around NaNs. In V3, finite running state is intentionally the right operand to the hardware max instruction; this preserves a valid prior maximum when one document lane produces NaN without introducing a generic standard-max comparison. A regression test covers this case.

Performance methodology

The primary evidence is the original strict three-way comparison of:

It used the official nine shapes, release builds, 50 measurements per shape, minimum ns/IP, and rotated implementation order to reduce ordering bias. Lower percentages are better. Reference is excluded from conclusions because this patch does not alter it. Auto duplicates V3 on this machine and is omitted below.

Aggregate results

ImplementationHEAD vs base medianHEAD vs base P90This PR vs base medianThis PR vs base P90This PR vs HEAD medianThis PR vs HEAD P90
f32 Scalar+1.96%+4.68%+0.35%+0.84%-1.27%-0.31%
f32 V3+0.73%+4.05%-0.20%0.00%-0.93%-0.45%
f16 V3+1.88%+60.60%0.00%+0.37%-2.46%+0.45%

The f16 V3 P90 for HEAD is dominated by one noisy/sensitive shape; the complete rows are included rather than hiding that variance.

Complete per-shape results

Type / ISAQDDimBase ns/IPHEAD ns/IPThis PR ns/IPHEAD vs baseThis PR vs baseThis PR vs HEAD
f32 Scalar16642566.83596.92386.8359+1.29%0.00%-1.27%
f32 Scalar3212501283.07253.15003.0950+2.52%+0.73%-1.75%
f32 Scalar3212838410.376010.461410.4126+0.82%+0.35%-0.47%
f32 Scalar32162566.84576.99226.8848+2.14%+0.57%-1.54%
f32 Scalar323251213.964814.082014.0039+0.84%+0.28%-0.55%
f32 Scalar64125051214.131214.143814.1000+0.09%-0.22%-0.31%
f32 Scalar64321283.07373.14943.0811+2.46%+0.24%-2.17%
f32 Scalar64322646.99227.12897.0508+1.96%+0.84%-1.10%
f32 Scalar8321283.17193.32033.1875+4.68%+0.49%-4.00%
f32 V316642564.15044.19924.1504+1.18%0.00%-1.16%
f32 V33212501282.06502.08002.0625+0.73%-0.12%-0.84%
f32 V3321283846.20126.23786.1890+0.59%-0.20%-0.78%
f32 V332162564.19924.27734.1895+1.86%-0.23%-2.05%
f32 V332325128.30088.35948.2812+0.71%-0.24%-0.93%
f32 V36412505128.32508.35008.3125+0.30%-0.15%-0.45%
f32 V364321282.08502.12402.0801+1.87%-0.23%-2.07%
f32 V364322644.27734.30664.2676+0.68%-0.23%-0.91%
f32 V38321284.24224.41414.2422+4.05%0.00%-3.89%
f16 V316642565.37115.42975.3613+1.09%-0.18%-1.26%
f16 V33212501282.24252.23502.2450-0.33%+0.11%+0.45%
f16 V3321283846.78716.95806.7871+2.52%0.00%-2.46%
f16 V332162566.75786.88486.6309+1.88%-1.88%-3.69%
f16 V3323251210.546912.343810.5859+17.04%+0.37%-14.24%
f16 V36412505128.55628.63128.5812+0.88%+0.29%-0.58%
f16 V364321282.80272.79302.7368-0.35%-2.35%-2.01%
f16 V364322647.138711.46487.0020+60.60%-1.92%-38.93%
f16 V38321286.12506.31256.1406+3.06%+0.26%-2.72%

Scope / non-goals

This PR intentionally does not add V4 or Neon kernels, change MinMax or Spherical behavior, modify ScopedAllocator, add cache detection, or introduce any other feature. It changes only the six MaxSim kernel paths needed to move reduction to the completed-panel boundary.

Validation

  • cargo fmt --all --check
  • cargo test -p diskann-quantization multi_vector (82 passed)
  • cargo clippy -p diskann-quantization --all-targets -- -D warnings
  • Full six-job benchmark input completed all nine shapes per job on the exact committed code

Limitations

This machine does not support AVX-512, so no V4 execution claim is made. Standalone later runs, including the exact-final completion run, experienced machine-frequency drift; they are used only to confirm that the exact code completes every benchmark case, not as cross-version percentage evidence. The strict rotated three-way table above remains the primary performance comparison.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
@partychen

Copy link
Copy Markdown
ContributorAuthor

AI: blocking: The benchmark section identifies the compared commits, shapes, measurement count, and ordering, but not the host hardware. For these ISA-sensitive results, please also record the CPU model, core topology, OS, and relevant frequency/power configuration so the measurements are reproducible and comparable (perf-claims-need-numbers).

b_row_bytes: NonZeroUsize,
b_rows: usize,
acc_bytes: usize,
work_bytes_per_b_row: usize,

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI: consider: work_bytes_per_b_row sounds like the complete working-set cost for one B row, but Plan::new multiplies it by AR. That makes the actual unit “bytes per A-row per B-row”; a future operation supplying an already-total per-B-row cost would over-reserve L1 by AR. Could we encode that unit in the parameter name (for example, work_bytes_per_a_row_per_b_row) and update the test wording? (names-match-vocabulary)

for chain in &chains[1..] {
for (m, c) in merged.iter_mut().zip(chain) {
*m = m.max_simd(*c);
// `run` seeds state with a finite identity, and every update preserves that invariant.

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI: nit: The invariant here is “non-NaN,” not “finite”: valid inputs or overflow can produce +∞, which remains a valid running maximum. Could this say that run seeds a non-NaN identity and that placing the prior non-NaN state on the right preserves it when a candidate is NaN?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@partychen