[mfoerste4]

Balanced k-means can return cluster centers at the origin. calc_centers_and_sizes divides each cluster's row sum by its size using raft::div_checkzero_op, so a cluster with no rows gets an all-zero center. Standard kmeans::fit avoids this because finalize_centroids copies the previous center back.

This can cause issues in downstream usage like IVF-PQ when all empty clusters have the same center -- a query close to the origin would then automatically probe all empty lists at once.

CC @tfeher, @huuanhhuyn

[coderabbitai]

No actionable comments were generated in the recent review.

Summary

Summary by CodeRabbit

• Bug Fixes
  • Empty clusters now retain their existing centers during incremental updates; newly initialized empty clusters keep zero-valued centers.
  • Incremental center calculations now account for existing cluster sizes, preventing centers from shifting unexpectedly when clusters receive no new points.
  • Balanced clustering now searches for eligible donor clusters more thoroughly when adjusting cluster sizes.

Balanced k-means now preserves centers for clusters with no points during center calculation. Initial cluster construction zero-initializes centers. New tests cover fitting with empty clusters and incremental updates.

Changes: Balanced k-means center handling

Layer / File(s)	Summary
Center accumulation and normalization cpp/src/cluster/detail/kmeans_balanced.cuh, cpp/src/cluster/kmeans_balanced.cuh, cpp/tests/cluster/kmeans_balanced.cu	Center calculation counts labels before reduction. Incremental updates weight existing centers by their prior cluster sizes. Empty clusters retain their centers. The API documentation describes when center values are read. The incremental update test checks that sizes become {4, 0} and both centers remain unchanged.
Initial center setup cpp/src/cluster/detail/kmeans_balanced.cuh, cpp/tests/cluster/kmeans_balanced.cu	build_clusters zero-initializes centers before calculating centers from initialized labels. The fit test checks that all center coordinates remain within the input data’s range.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: bdice, cjnolet, tarang-jain

Merge Risk: Minimal · up to c654a

Balanced k-means now keeps the centers of empty clusters instead of moving them to the origin. No actionable merge-blocking risk was found in the supplied changes.

Architecture risk: Low · up to c654a

The change affects 1 system.

Changed systems: cpp

Architecture concerns
No architecture-level concerns identified.

Review details

Systems and components

• observed — cpp (service) was modified; 3 changed files map to changed impact.

Before / after behavior

• observed — Modified behavior in cpp/src/cluster/kmeans_balanced.cuh: The documentation adds that a cluster with no points and zero prior weight retains its center.
• observed — Modified behavior in cpp/src/cluster/kmeans_balanced.cuh: The parameter documentation changes centroids from output-only to input/output and states that existing centers are read when reset_counters is false or a cluster has no points.
• observed — Modified behavior in cpp/tests/cluster/kmeans_balanced.cu: Adds a fit test using four repeated row patterns and 16 clusters; it checks that all center coordinates remain between 0.999 and 4.001.
• observed — Modified behavior in cpp/tests/cluster/kmeans_balanced.cu: Adds an incremental center-update test with four rows assigned to one of two clusters. It checks sizes become {4, 0} and that both centers, including the empty cluster’s center, remain unchanged.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title names the bug and its subject: balanced k-means moves centers of empty clusters to the origin. The changeset fixes this in calc_centers_and_sizes, build_clusters, and the random donor logic,…
Description check		The description states the problem: balanced k-means returns zero centers for empty clusters because calc_centers_and_sizes uses raft::div_checkzero_op. It contrasts this with kmeans::fit, which copie…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Title check

Explanation

The title names the bug and its subject: balanced k-means moves centers of empty clusters to the origin. The changeset fixes this in calc_centers_and_sizes, build_clusters, and the random donor logic, and adds tests for empty-cluster center retention. The title is specific and matches the main change.

Full details: Description check

Explanation

The description states the problem: balanced k-means returns zero centers for empty clusters because calc_centers_and_sizes uses raft::div_checkzero_op. It contrasts this with kmeans::fit, which copies the previous center back, and gives a downstream IVF-PQ impact. This matches the centroid-handling changes and the new tests in the changeset.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/cuvs
• Branch: main
• File: .coderabbit.yaml
• Commit: 471e6d7
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: da5ca913-f687-4308-a980-cb0b9e751f25

Commits

Reviewing files that changed from the base of the PR and between 04f504b and c654a36.

Files selected for processing: 1

• cpp/src/cluster/detail/kmeans_balanced.cuh

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

[huuanhhuyn]

Should we have a small comment here to explain constant value of each row?
Similarly at the assertion line?