Skip to content

Add support for PQ preprocessing API #1278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
rapids-bot merged 90 commits into from
Jan 23, 2026
Merged

Add support for PQ preprocessing API #1278

rapids-bot merged 90 commits into from
Jan 23, 2026

Conversation

@lowener
Copy link
Contributor

@lowener lowener commented Aug 23, 2025 •
edited by cjnolet
Loading

Closes #107

This PR adds support for a PQ preprocessing API. It gives access to train(), transform() and inverse_transform() function that can be used to transform a dataset into PQ codes. It is re-using the VPQ functions from CAGRA-Q.

@lowener
Initial commit for PQ preprocessing API
cfe4f92 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@copy-pr-bot
Copy link

copy-pr-bot bot commented Aug 23, 2025

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@lowener lowener added feature request New feature or request non-breaking Introduces a non-breaking change C++ labels Aug 25, 2025
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
cjnolet
cjnolet reviewed Aug 25, 2025
View reviewed changes
@lowener
Support n_lists and cleanup code
096daa5 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@cjnolet cjnolet moved this from Todo to In Progress in Vector Search, ML, & Data Mining Release Board Sep 5, 2025
gland1
gland1 reviewed Sep 14, 2025
View reviewed changes
gland1
gland1 reviewed Sep 14, 2025
View reviewed changes
@lowener
Switch to VPQ
244a9cd 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Fix trainpq and train workflow
9537eb1 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Remove timer
2883a25 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Merge branch 'branch-25.10' into 25.10-pq-preprocessing
78dfd69
@lowener
Cleanup Code
9dd0cfe 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Add double dtype
9c543fb 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener lowener marked this pull request as ready for review September 24, 2025 16:09
@lowener lowener requested review from a team as code owners September 24, 2025 16:09
KyleFromNVIDIA
KyleFromNVIDIA approved these changes Sep 24, 2025
View reviewed changes
Copy link
Member

@KyleFromNVIDIA KyleFromNVIDIA left a comment

Choose a reason for hiding this comment

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

Approved trivial CMake changes

@lowener lowener requested review from a team as code owners January 23, 2026 14:29
@lowener lowener requested a review from AyodeAwe January 23, 2026 14:29
robertmaynard
robertmaynard requested changes Jan 23, 2026
View reviewed changes
Copy link
Contributor

@robertmaynard robertmaynard left a comment

Choose a reason for hiding this comment

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

This PR pulls CI changes from main into release/26.02 which isn't allowed.

You need to rewrite your merge history and remove these changes.

@irina-resh-nvda
Copy link
Contributor

irina-resh-nvda commented Jan 23, 2026

Thank you @lowener for this PR! parse_build_param function on line 198 in /cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h also needs to pick up the new parameters

robertmaynard
robertmaynard approved these changes Jan 23, 2026
View reviewed changes
Copy link
Contributor

@robertmaynard robertmaynard left a comment

Choose a reason for hiding this comment

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

Aprroving as the 26.04 CI changes have been removed, thanks

aamijar and others added 12 commits January 23, 2026 09:08
@aamijar @lowener
Spectral Embedding with all_neighbors (rapidsai#1693)
8052db3 
Instead of using `cuvs::neighbors::brute_force::build` and `cuvs::neighbors::brute_force::search` we can consolidate to use the `cuvs::neighbors::all_neighbors::build` API.

Authors:
  - Anupam (https://github.com/aamijar)

Approvers:
  - Tarang Jain (https://github.com/tarang-jain)
  - Jinsol Park (https://github.com/jinsolp)
  - Victor Lafargue (https://github.com/viclafargue)

URL: rapidsai#1693
@jinsolp @lowener
Deduplicate calc_chunk_indices_kernel (rapidsai#1657)
c927dec 
Closes rapidsai#1577

Reduces binary size by deduplicating `calc_chunk_indices_kernel`. This PR reduces instantiations from 62 -> 1 for each template (`BlockDim`=32, 64, ..., 1024)

### Binary Size Changes
CUDA 12.9: 1096.15MB -> 
CUDA 13: 432.98 MB->

Authors:
  - Jinsol Park (https://github.com/jinsolp)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - Robert Maynard (https://github.com/robertmaynard)

URL: rapidsai#1657
@AyodeAwe @lowener
Prepare release/26.02
60ebe3f
@mmccarty @lowener
wheel builds: react to changes in pip's handling of build constraints (…
f198537 
…rapidsai#1710)

Contributes to rapidsai/build-planning#242

Modifying `ci/build_wheel.sh` to

- pass`--build-constraint="${PIP_CONSTRAINT}"` unless build isolation is enabled.
- unset `PIP_CONSTRAINT` (set by rapids-init-pip)... it doesn't affect builds as of pip 25.3, and results in an error from `pip wheel` when set and `--build-constraint` is also passed

Authors:
  - Mike McCarty (https://github.com/mmccarty)

Approvers:
  - James Lamb (https://github.com/jameslamb)

URL: rapidsai#1710
@lowener
Use raft::TxN_t
4bb9435 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Use separate vector for VQ labels, switch Vamana to public PQ API
2e13085 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Add note on doc
081254c 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Add issue #
a69e2cf 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener
Fix doc
6e77ab0 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@robertmaynard @lowener
pre-built libcuvs_c.so now use the new ABI major/minor values (rapids…
841a3a9 
…ai#1708)

Based on the in-development cuVS ABI stability docs we now encode the C API's ABI stability in the SOVERSION

Authors:
  - Robert Maynard (https://github.com/robertmaynard)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: rapidsai#1708
@robertmaynard @lowener
Correct base release for cuvs abi 1 major (rapidsai#1724)
b95eb46 
Correct the libcuvs_c.so ABI major to start at 26.02

Authors:
  - Robert Maynard (https://github.com/robertmaynard)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: rapidsai#1724
@lowener
Add new option to ann-bench
809bacb 
Signed-off-by: Mickael Ide <mide@nvidia.com>
@lowener lowener force-pushed the 25.10-pq-preprocessing branch from 70a2564 to 809bacb Compare January 23, 2026 17:11
@mythrocks
Copy link
Contributor

mythrocks commented Jan 23, 2026 •
edited
Loading

I don't see changes to the Java code here, so I'm reviewing these from the ABI perspective.

Will we be addressing this comment about enum names, in this PR/release? I ask because I didn't see a PR attached to #1717.

@cjnolet cjnolet removed request for a team and AyodeAwe January 23, 2026 19:08
@lowener
Copy link
Contributor Author

lowener commented Jan 23, 2026 •
edited
Loading

@mythrocks A force push triggered a request for re-review on all files.
The comment about ABI changes was addressed, and now the enums introduced in this PR have a name more specific that is not likely to cause symbol collision.
The issue #1717 is to fix it accross all the codebase but it is not within the scope of this PR

@cjnolet
Copy link
Member

cjnolet commented Jan 23, 2026

/merge

cjnolet
cjnolet approved these changes Jan 23, 2026
View reviewed changes
cpp/include/cuvs/preprocessing/quantize/pq.hpp
}

/**
* @brief Applies inverse quantization transform to given dataset
Copy link
Member

@cjnolet cjnolet Jan 6, 2026

Choose a reason for hiding this comment

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

Suggested change
* @brief Applies inverse quantization transform to given dataset
* @brief Applies inverse quantization transform to given set of encoded vectors

c/include/cuvs/cluster/kmeans.h Outdated
/**
* @brief Type of k-means algorithm.
*/
typedef enum { KMeans = 0, KMeansBalanced = 1 } cuvsKMeansType;
Copy link
Member

@cjnolet cjnolet Jan 6, 2026

Choose a reason for hiding this comment

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

Future task- we should update cuvsKMeansParams_t to use this and remove the bool hierarchical flag. Can you create an issue for this just so we don't forget it? It doesn't have to be done now, especially since we're striving to maintain ABI compatibility in our C APIs (we need to start deprecating breaking changes).

@rapids-bot rapids-bot bot merged commit d1b11d5 into rapidsai:release/26.02 Jan 23, 2026
191 of 193 checks passed
rapids-bot bot pushed a commit that referenced this pull request Jan 24, 2026
@tarang-jain
[FEA] IVF-PQ to Write Flat PQ Codes (#1607)
d291569 
This PR brings new params to ivf_pq: an option for the user to choose the layout of the ivf lists. The lists can be flat (no interleaving) or interleaved (current default). Flat codes allows building the index in a CPU-compatible format.

[UPDATE as of 12/19/2025]:
After #1278 is merged, we can unify IVF-PQ and PQ API codepaths.

[UPDATE 01/08/2026]:
This PR can be merged before #1278. The flat code-writing can potentially be reverted once #1278 is merged (so we can later use the PQ preprocessing API directly). However that will come naturally as a part if a broader unification of IVF-PQ and PQ codepaths.

[Benchmarks 01/15/2026]:
## IVF-PQ Layout Benchmark Results

**Dataset**: 1,000,000 vectors × 128 dimensions | **pq_dim**: 32

pq_bits | Code Size | Direct FLAT Build (ms) | INTERLEAVED Build (ms) | Convert INTERLEAVED to FLAT with Codepacker (ms) | Total time for INTERLEAVED build + Conversion to FLAT with Codepacker (unpack) (ms) | Overhead |
|:-------:|:---------:|:---------------:|:----------------------:|:-------------------:|:----------------------:|:--------:|
| 8 | 32 bytes | 372.46 | 385.86 | 985.28 | 1371.13 | 3.68× |
| 6 | 24 bytes | 298.83 | 300.99 | 961.82 | 1262.82 | 4.23× |
| 5 | 20 bytes | 283.25 | 281.95 | 795.43 | 1077.38 | 3.80× |
| 4 | 16 bytes | 270.63 | 271.01 | 489.73 | 760.75 | 2.81× |

Authors:
  - Tarang Jain (https://github.com/tarang-jain)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)
  - Robert Maynard (https://github.com/robertmaynard)

URL: #1607
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@benfred benfred benfred left review comments

@divyegala divyegala divyegala left review comments

@tarang-jain tarang-jain tarang-jain left review comments

@KyleFromNVIDIA KyleFromNVIDIA KyleFromNVIDIA approved these changes

@robertmaynard robertmaynard robertmaynard approved these changes

@cjnolet cjnolet cjnolet approved these changes

@tfeher tfeher tfeher approved these changes

+2 more reviewers

@gland1 gland1 gland1 left review comments

@mfoerste4 mfoerste4 mfoerste4 left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

@lowener lowener

Labels

C++ feature request New feature or request non-breaking Introduces a non-breaking change

Projects

Vector Search, ML, & Data Mining Release Board
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.