Skip to content

PiPNN 7/7: add vertical partition Top-K - #1366

Draft
weiyaoluo (SeliMeli) wants to merge 228 commits into
pipnn-stack/06-hash-prunefrom
pipnn-stack/07-vertical-topk
Draft

PiPNN 7/7: add vertical partition Top-K#1366
weiyaoluo (SeliMeli) wants to merge 228 commits into
pipnn-stack/06-hash-prunefrom
pipnn-stack/07-vertical-topk

Conversation

@SeliMeli

@SeliMeli weiyaoluo (SeliMeli) commented Sep 1, 2026

Copy link
Copy Markdown
Contributor
  • Does this PR have a descriptive title that could go in our release notes?
  • Does this PR add any new dependencies? No.
  • Does this PR modify any existing APIs? No public API changes.
  • Is the change to the API backwards compatible? Yes.
  • Should this result in any changes to our documentation, either updating existing docs or adding new ones? No.

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 f32 leader IDs.

Data flow

  1. Partition GEMM writes leader-major dots as [leader, point].
  2. A complete SIMD group loads one leader's scores across point lanes.
  3. Runtime rank vectors keep score and exact u32 ID state lane-for-lane.
  4. Reusable worker scratch owns the O(K × SIMD-width) vectors and row-major ID scatter buffer.
  5. A scalar tail handles points after the complete SIMD prefix.

The metric seam supplies matching vertical formulas for L2, Cosine, normalized Cosine, and inner product.

Invariants

  • Runtime K has no fixed ceiling.
  • Strict comparisons preserve sampled-leader order for equal scores.
  • NaN and positive infinity do not displace rankable candidates.
  • Finite f32::MAX remains rankable.
  • Unfilled ranks retain u32::MAX and never become indexed leader IDs.
  • Score and exact-ID vectors use the same compile-time lane policy.
  • No per-point allocation occurs; worker scratch grows and is reused.

Review order

  1. diskann/src/graph/pipnn/simd.rs: score-lane to exact-ID-vector mapping and mask conversion.
  2. diskann/src/graph/pipnn/kernel_metric/partition.rs: leader-major metric formulas.
  3. diskann/src/graph/pipnn/partition_kernel.rs: vertical insertion, ID scatter, and scalar tail.
  4. 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.
  • Temporary one-line f32x16 to f32x8 policy switch — 40 focused tests passed; source restored.
  • Native V4 focused run — 40 passed.
  • cargo clippy -p diskann --features pipnn --all-targets -- -D warnings — passed.
  • AArch64 and Windows GNU cross-target 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.

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.
@SeliMeli weiyaoluo (SeliMeli) changed the title pipnn stack/07 vertical topk PiPNN 7/7: add vertical partition Top-K Sep 1, 2026
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.
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/07-vertical-topk branch from fb227e2 to a695f91 Compare September 2, 2026 08:45
Sign up for free to 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