[WIP] IP Vector Normalization to avoid all vectors clumped into single cluster in IVF-PQ by HowardHuang1 · Pull Request #1892 · rapidsai/cuvs

HowardHuang1 · 2026-03-07T06:42:47Z

Addresses #1875.

Current Implementation: Have Inner Product code path match Cosine code path.
Within kmeans, the E-step of cluster assignment will be operating on normalized vectors but the M-step recomputes centers based on raw vectors. Search still happens on raw vectors.

In each EM iteration, before predict() we L2 row-normalize the centroid matrix (cluster_centroids) for IP. Within the E-step's call to predict(), we compute the norm on all vectors in the dataset and dataset_norm is passed into predict_core() and used in fusedDistanceNNMinReduce<> kernel. This is a norm to norm comparison between normalized vectors and normalized centroids.

Then in the M-Step, calc_centers_and_sizes() recalculates the centroids based on the raw dataset and rewrites raw centroids into cluster_centers.

For the code review, the bulk of the changes are in this file: ‎cpp/src/cluster/detail/kmeans_balanced.cuh
The rest have to do with linking and can be ignored. Those files will be removed prior to merging.

So far, we've attempted:

Normalization of whole IVF-PQ path by normalizing input dataset itself
Normalization of kmeans only by making a copy (2x memory spike) and recovering raw vectors for search
[**Current Implementation ] Setting IP to same code path as Cosine for normalization.
[**In progress ] Normalization within kernels themselves.

Pareto plots for approach #3 (working):
IVF-PQ only:

Cagra end-to-end (Showing no harm in search performance):
Showing Fix for Oracle's Issue (in Cagra):

… cases due to reduced precision 8-bit values stored in LUT

copy-pr-bot · 2026-03-07T06:42:50Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

aamijar · 2026-03-09T22:59:53Z

/ok to test 78e47f1

aamijar · 2026-03-09T23:05:27Z

Hi @HowardHuang1, are we able to document an example in the description of the PR where vectors were clumped in a single cluster, and then verify that the fix (this PR) solves that case?

Also you can run pre-commit to fix the CI style check.

HowardHuang1 · 2026-03-09T23:15:11Z

Yes will do. I'll upload a before and after example that illustrates the fixed cluster distributions and timing improvements for search.

…d IVF-PQ IP

aamijar · 2026-03-10T22:15:48Z

/ok to test e822460

…reviously for debug only

aamijar · 2026-03-14T00:23:36Z

/ok to test da32073

…ange to only apply normalization to kmeans step

….s. raw doesn't affect ordering of clusters for nearest cluster assignment

aamijar · 2026-03-18T22:33:19Z

Hi @HowardHuang1, is this targeting release/26.04? If so, please retarget from main to release/26.04

jinsolp · 2026-03-19T02:51:14Z

Hi @HowardHuang1 you'll have to rebase on release/26.04 if you want to target that branch. 🙂

…put command parsing

jinsolp

Thanks @HowardHuang1 ! I have a few questions as I try to understand the fix in this code.

So we normalize the data, and train kmeans on this normalized data. However, instead of saving the resulting centroids (which come from the normalized space), we re-compute the proper centroids on the raw vectors (what calc_centers_and_sizes do)?
what are the changes in the cuvs bench for? I'm not so familiar with this part of the code so this is out of curiosity.
If I'm reading things right in the screenshot, we get a lower recall with the normalized vectors compared to using the raw vectors? And the search also takes longer?

jinsolp · 2026-03-20T22:25:16Z

+        handle, kmeans_params, trainset_kmeans_const_view, centers_const_view, labels_view);
+      // Recompute centers in original space (mean of unnormalized trainset per cluster), overwrites centers_view
+      rmm::device_uvector<uint32_t> cluster_sizes(impl->n_lists(), stream, device_memory);
+      cuvs::cluster::kmeans::detail::calc_centers_and_sizes<float, float, internal_extents_t, uint32_t, uint32_t>(


can we use the calc_centers_and_sizes in the public namespace if possible?

Hey @jinsolp !

Yes exactly. If we were to keep the normalized centroids that would cause the computation to no longer be Inner Product and instead it will degenerate to Cosine which is not what we want. To prevent this, we use calc_centers_and_sizes to re-compute the proper centroids on the raw vectors. In essence, we only want normalization to happen in the kmeans cluster assignment step rather than the whole pipeline (normalization of the whole pipeline changes metric to Cosine).

The changes in cuvs bench are to resolve a linker issue. Will remove these when merging. They can be ignored for now.

The screenshot is a bit outdated, will replace soon with a new one. Recall should be same search should be faster.

…inner product

cjnolet · 2026-03-26T00:39:56Z

+    if (impl->metric() == distance::DistanceType::InnerProduct) {
+      // Normalization only for k-means: use a copy so trainset stays in original space; metric
+      // remains inner product for the rest of the pipeline.
+      auto trainset_kmeans = raft::make_device_mdarray<float>(


Wait- we seem to be doing a copy of the trainset? Why? Just do the normalization in the distance kernels.

This is not the way we do things- we don't do unecessary copies just to avoid changing kernels. This becomes frustrating for users, because they are almost always memory limited and every GB counts.

…ulting in 2x memory spike

…ng fused cosine implementation

…ld is called because only ivf_pq should have normalization for IP, other index types should not --> this version fixes low recall issue on gist and fixes Oracle bug

Fixed vector normalization for Inner Product -- still failing 15 test…

ee94971

… cases due to reduced precision 8-bit values stored in LUT

HowardHuang1 requested review from a team as code owners March 7, 2026 06:42

github-project-automation bot added this to Unstructured Data Processing Mar 7, 2026

aamijar assigned HowardHuang1 Mar 9, 2026

aamijar moved this to In Progress in Unstructured Data Processing Mar 9, 2026

aamijar added non-breaking Introduces a non-breaking change bug Something isn't working labels Mar 9, 2026

Merge branch 'main' into HH-Vector-Normalization

78e47f1

add logging to compare recall and search speed for raw v.s. normalize…

e822460

…d IVF-PQ IP

revert change comparing normalize and raw in single run -- that was p…

da32073

…reviously for debug only

HowardHuang1 added 5 commits March 16, 2026 15:37

previously normalization was applied to entire IVF-PQ pipeline --> ch…

5375bb5

…ange to only apply normalization to kmeans step

revert to raw vectors. No need to normalize here because normalized v…

c320640

….s. raw doesn't affect ordering of clusters for nearest cluster assignment

clean up code

0486bf5

clean up code

107e2b3

upload code that resolves linker issue + live csv updates

5442b89

HowardHuang1 requested a review from a team as a code owner March 18, 2026 16:40

remove live_csv

dc9b6df

aamijar removed the request for review from a team March 18, 2026 22:32

Merge branch 'main' into HH-Vector-Normalization

cf3666c

HowardHuang1 changed the base branch from main to release/26.04 March 18, 2026 22:41

HowardHuang1 requested review from a team as code owners March 18, 2026 22:41

HowardHuang1 requested a review from AyodeAwe March 18, 2026 22:41

HowardHuang1 changed the base branch from release/26.04 to main March 18, 2026 23:41

jinsolp removed request for a team and AyodeAwe March 19, 2026 02:51

HowardHuang1 added 6 commits March 18, 2026 21:07

hardcode file path instead of searching multiple directories + fix in…

bdda881

…put command parsing

clean up unnecessary checks in data_export.py

c237db3

bring back comma parsing instead of underscore parsing

3c20377

bring back parts of plot/__main__.py for clarity

728c964

get rid of incremental JSON->CSV write for clarity

6c9bc36

bring back original plot/__main__.py for clarity

eb6bb88

jinsolp reviewed Mar 20, 2026

View reviewed changes

fix cuvs-bench generate groundtruth which was sorted incorrectly for …

4546a85

…inner product

cjnolet reviewed Mar 26, 2026

View reviewed changes

HowardHuang1 changed the base branch from main to release/26.04 March 30, 2026 18:37

HowardHuang1 changed the base branch from release/26.04 to main March 30, 2026 18:38

HowardHuang1 added 3 commits March 30, 2026 15:43

revert normalization outside kernel that requires copy of dataset res…

ea8a6b4

…ulting in 2x memory spike

didn't modify kernels themselves but routed IP predict through existi…

5800e53

…ng fused cosine implementation

add inner_product_cosine_assignment flag which is set when ivf_pq_bui…

945d3ee

…ld is called because only ivf_pq should have normalization for IP, other index types should not --> this version fixes low recall issue on gist and fixes Oracle bug

Conversation

HowardHuang1 commented Mar 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot bot commented Mar 7, 2026

Uh oh!

aamijar commented Mar 9, 2026

Uh oh!

aamijar commented Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HowardHuang1 commented Mar 9, 2026

Uh oh!

aamijar commented Mar 10, 2026

Uh oh!

aamijar commented Mar 14, 2026

Uh oh!

aamijar commented Mar 18, 2026

Uh oh!

jinsolp commented Mar 19, 2026

Uh oh!

jinsolp left a comment

Choose a reason for hiding this comment

Uh oh!

jinsolp Mar 20, 2026

Choose a reason for hiding this comment

Uh oh!

HowardHuang1 Mar 20, 2026

Choose a reason for hiding this comment

Uh oh!

cjnolet Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

cjnolet Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HowardHuang1 Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

HowardHuang1 commented Mar 7, 2026 •

edited

Loading

aamijar commented Mar 9, 2026 •

edited

Loading

cjnolet Mar 26, 2026 •

edited

Loading