PiPNN 7/7: add vertical partition Top-K - #1366
Draft
weiyaoluo (SeliMeli) wants to merge 228 commits into
Draft
Conversation
Select architecture, metric, and leaf-width implementations once, then reuse direct diskann-wide function pointers across stripes and leaves. BREAKING CHANGE: callers construct LeafKernel or PartitionKernel and pass MatrixView-backed inputs and outputs.
Use output columns as the sole leaf-specific neighbor count and reserve row/column terminology for matrix shapes. BREAKING CHANGE: LeafKernel::new no longer takes k, nearest_neighbors returns (), and kernel input/neighbor/error fields use source-target and point-leader names.
Keep PiPNN beside graph policy so later layers can reuse private RobustPrune state without publishing it across a crate boundary. Preserve independent kernel oracles while removing duplicate formula-sharing differential wrappers.
Remove the submitted DiskANN microbenchmark target and co-locate numerical tests with their implementation files.
Add only the PiPNN feature and SDE package coverage; leave workflow triggers and existing formatting unchanged.
Remove repeated CSR and point-range validation from the private leaf path. Keep checked platform slab allocation boundaries.
Keep exact candidate, hash, conversion, and stale-scratch behavior checks while dropping layout-only assertions.
Replace opaque generated inputs with named geometry, split distinct output contracts, and parameterize validation boundaries.
Make zero-k, zero-point, and zero-dimension conditions explicit in test fixtures.
Reuse the repository nightly Miri job and keep strict-provenance checks limited to exact pointer-boundary tests.
Store leader-major dots and maintain one runtime-width Top-K per SIMD point lane. Keep score and exact u32 ID state in reusable worker scratch. Use scalar tails for partial point groups.
weiyaoluo (SeliMeli)
force-pushed
the
pipnn-stack/07-vertical-topk
branch
from
September 2, 2026 08:45
fb227e2 to
a695f91
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Reference Issues/PRs
Stack parent: #1295
What does this implement/fix? Briefly explain your changes.
Purpose
Replace Partition's leader-lane ranking with a vertical Top-K kernel. Each SIMD lane owns one point and maintains an independent runtime-width Top-K while leaders are scanned in sampled order.
This PR does not include the fixed-K POC paths, feature toggles, benchmark code, Leaf changes, stripe tuning, or
f32leader IDs.Data flow
[leader, point].u32ID state lane-for-lane.The metric seam supplies matching vertical formulas for L2, Cosine, normalized Cosine, and inner product.
Invariants
f32::MAXremains rankable.u32::MAXand never become indexed leader IDs.Review order
diskann/src/graph/pipnn/simd.rs: score-lane to exact-ID-vector mapping and mask conversion.diskann/src/graph/pipnn/kernel_metric/partition.rs: leader-major metric formulas.diskann/src/graph/pipnn/partition_kernel.rs: vertical insertion, ID scatter, and scalar tail.diskann/src/graph/pipnn/partitioning.rs: architecture-typed worker scratch ownership.Validation
cargo test -p diskann --features pipnn graph::pipnn::partition_kernel— 40 passed.cargo test -p diskann --features pipnn— 649 passed; doctests passed.f32x16tof32x8policy switch — 40 focused tests passed; source restored.cargo clippy -p diskann --features pipnn --all-targets -- -D warnings— passed.cargo check— passed.cargo fmt --all --check,git diff --check, LSP diagnostics, and lens diagnostics — passed.No performance claim is made in this PR. Representative end-to-end profiling remains a follow-up gate before claiming a build-time improvement.