Repository navigation
Implemantion for two-level KMeans Trees in ScaNN - #1878
Merged
Merged
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
rmaschal
changed the title
rmaschal
added
feature request
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds optional coarse-level clustering to the ScaNN index (API and device storage), introduces batched k‑means train/predict helpers, integrates coarse clustering and centers_view into the build/quantization flow, extends serialization to persist coarse artifacts, and adds a notebook to emit ScaNN protobuf/assets. ChangesCoarse Clustering Support
Protobuf Example Notebook
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Ruff (0.15.12)notebooks/scann_protobuf_example.ipynbUnexpected end of JSON input Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (5)
cpp/src/neighbors/scann/detail/scann_serialize.cuh (1)
77-107: 💤 Low valueRemove commented-out code.
Lines 84-85 contain leftover commented-out code that appears to be from the previous implementation. This should be cleaned up.
Proposed fix
void save_labels(raft::resources const& res, std::filesystem::path labels_path, raft::device_vector_view<const uint32_t, IdxT> labels_view, raft::device_vector_view<const uint32_t, IdxT> soar_labels_view) { auto combined_labels = raft::make_device_vector<int, IdxT>(res, 2 * labels_view.extent(0)); - // auto labels_view = index_.labels(); - // auto soar_labels_view = index_.soar_labels(); raft::linalg::map_offset(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/scann/detail/scann_serialize.cuh` around lines 77 - 107, In save_labels remove the leftover commented-out lines "// auto labels_view = index_.labels();" and "// auto soar_labels_view = index_.soar_labels();" since labels_view and soar_labels_view are already parameters; simply delete those two commented lines near the creation of combined_labels to clean up the function and avoid stale code.cpp/include/cuvs/neighbors/scann.hpp (1)
178-196: 💤 Low valueConsider adding Doxygen comments to new public accessors.
The new
coarse_centers(),coarse_labels(), andcoarse_soar_labels()accessors lack Doxygen documentation. While this is consistent with the existing accessors (centers(),labels(), etc.), the coding guidelines recommend complete Doxygen documentation for public API functions. Consider adding brief@briefcomments to improve documentation coverage.Also applies to: 219-244, 304-306
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/cuvs/neighbors/scann.hpp` around lines 178 - 196, The new public accessor methods coarse_centers(), coarse_labels(), and coarse_soar_labels() are missing Doxygen documentation; add brief Doxygen comments (e.g., a one-line `@brief` describing what each accessor returns and any important details such as tensor shape/semantic) above each accessor to match existing style used for centers()/labels(), ensuring documentation coverage for the public API; apply the same `@brief` comment pattern to the other new accessors noted in the diff (the other similar accessor blocks) so all public accessors are documented consistently.cpp/src/neighbors/scann/detail/scann_build.cuh (2)
114-117: 💤 Low valueRedundant
sync_streamcall.Line 114 already synchronizes inside the
elsebranch, making the unconditional sync at line 116 redundant for that path. Consider removing line 114 since line 116 handles synchronization for both branches.Proposed fix
cuvs::cluster::kmeans::fit(res, kmeans_params, raft::make_const_mdspan(trainset.view()), std::nullopt, centers, raft::make_host_scalar_view<float>(&inertia), raft::make_host_scalar_view<int64_t>(&n_iter)); - - raft::resource::sync_stream(res); } raft::resource::sync_stream(res); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/scann/detail/scann_build.cuh` around lines 114 - 117, The duplicate synchronization call to raft::resource::sync_stream(res) is redundant: remove the earlier call inside the conditional branch so only the final unconditional raft::resource::sync_stream(res) remains; locate the redundant call where sync_stream(res) appears just before the closing brace of the else/conditional block (referencing the res variable and raft::resource::sync_stream) and delete that inner invocation so the single final sync covers both branches.
39-52: 💤 Low valueMinor typos in Doxygen comments.
The parameter documentation contains typos:
- Line 40: "balaned" → "balanced"
- Lines 43, 123: "
@tparms" → "@tparam"- Lines 43, 123: "IxT" → "IdxT"
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/scann/detail/scann_build.cuh` around lines 39 - 52, Fix the typos in the Doxygen block for the kmeans training helper: change "balaned" to "balanced" in the description, correct both occurrences of the malformed tag "`@tparms`" to "`@tparam`", and rename the template parameter label "IxT" to "IdxT" so the template docs match the actual template naming (look for the Doxygen block that documents template parameters T and IxT/IdxT for the kmeans centers training helper).notebooks/scann_protobuf_example.ipynb (1)
412-421: ⚡ Quick winRedundant assignment and unused parameter.
- Line 413 sets
bfloat16.enabled = params.bf16, then lines 415-416 redundantly set it again inside theif params.bf16block.- Line 417 hardcodes
noise_shaping_threshold = 0.2instead of usingparams.bf16_threshold.Proposed fix
" scann_config.exact_reordering.approx_num_neighbors = 100\n", - " scann_config.exact_reordering.bfloat16.enabled = params.bf16\n", - "\n", " if params.bf16:\n", - " scann_config.exact_reordering.bfloat16.enabled = params.bf16\n", - " scann_config.exact_reordering.bfloat16.noise_shaping_threshold = 0.2\n", + " scann_config.exact_reordering.bfloat16.enabled = True\n", + " scann_config.exact_reordering.bfloat16.noise_shaping_threshold = params.bf16_threshold\n",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@notebooks/scann_protobuf_example.ipynb` around lines 412 - 421, Remove the redundant unconditional assignment to scann_config.exact_reordering.bfloat16.enabled and instead only set bfloat16.enabled and bfloat16.noise_shaping_threshold inside the guarded branch when params.bf16 is truthy; replace the hardcoded noise_shaping_threshold value 0.2 with params.bf16_threshold so scann_config is built from the params object (update the block that touches scann_config, params, and scann_config_path accordingly).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/include/cuvs/neighbors/scann.hpp`:
- Around line 178-196: The new public accessor methods coarse_centers(),
coarse_labels(), and coarse_soar_labels() are missing Doxygen documentation; add
brief Doxygen comments (e.g., a one-line `@brief` describing what each accessor
returns and any important details such as tensor shape/semantic) above each
accessor to match existing style used for centers()/labels(), ensuring
documentation coverage for the public API; apply the same `@brief` comment pattern
to the other new accessors noted in the diff (the other similar accessor blocks)
so all public accessors are documented consistently.
In `@cpp/src/neighbors/scann/detail/scann_build.cuh`:
- Around line 114-117: The duplicate synchronization call to
raft::resource::sync_stream(res) is redundant: remove the earlier call inside
the conditional branch so only the final unconditional
raft::resource::sync_stream(res) remains; locate the redundant call where
sync_stream(res) appears just before the closing brace of the else/conditional
block (referencing the res variable and raft::resource::sync_stream) and delete
that inner invocation so the single final sync covers both branches.
- Around line 39-52: Fix the typos in the Doxygen block for the kmeans training
helper: change "balaned" to "balanced" in the description, correct both
occurrences of the malformed tag "`@tparms`" to "`@tparam`", and rename the template
parameter label "IxT" to "IdxT" so the template docs match the actual template
naming (look for the Doxygen block that documents template parameters T and
IxT/IdxT for the kmeans centers training helper).
In `@cpp/src/neighbors/scann/detail/scann_serialize.cuh`:
- Around line 77-107: In save_labels remove the leftover commented-out lines "//
auto labels_view = index_.labels();" and "// auto soar_labels_view =
index_.soar_labels();" since labels_view and soar_labels_view are already
parameters; simply delete those two commented lines near the creation of
combined_labels to clean up the function and avoid stale code.
In `@notebooks/scann_protobuf_example.ipynb`:
- Around line 412-421: Remove the redundant unconditional assignment to
scann_config.exact_reordering.bfloat16.enabled and instead only set
bfloat16.enabled and bfloat16.noise_shaping_threshold inside the guarded branch
when params.bf16 is truthy; replace the hardcoded noise_shaping_threshold value
0.2 with params.bf16_threshold so scann_config is built from the params object
(update the block that touches scann_config, params, and scann_config_path
accordingly).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1c2ef832-dfa2-4c6f-af2b-0075349d99c8
📒 Files selected for processing (5)
cpp/include/cuvs/neighbors/scann.hppcpp/src/neighbors/scann/detail/scann_avq.cuhcpp/src/neighbors/scann/detail/scann_build.cuhcpp/src/neighbors/scann/detail/scann_serialize.cuhnotebooks/scann_protobuf_example.ipynb
💤 Files with no reviewable changes (1)
- cpp/src/neighbors/scann/detail/scann_avq.cuh
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/include/cuvs/neighbors/scann.hpp (2)
178-196:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMissing validation for invalid coarse-cluster parameter combinations.
params.n_coarse_clustersis forwarded but never validated. At minimum, guard invalid combinations (e.g.,n_coarse_clusters > 0 && n_coarse_clusters > n_leaves) before allocations/training paths.Proposed guard
{ RAFT_EXPECTS(params.pq_bits == 4 || params.pq_bits == 8, "ScaNN only supports 4 or 8 bit PQ"); RAFT_EXPECTS(dim >= params.pq_dim, "PQ subspace dimension (pq_dim) should be smaller than the dataset dimension"); RAFT_EXPECTS(dim % params.pq_dim == 0, "PQ subspace dimension (pq_dim) must divide the dataset dimension"); + RAFT_EXPECTS(params.n_coarse_clusters <= params.n_leaves, + "n_coarse_clusters must be <= n_leaves"); }As per coding guidelines, "Input validation must check ... invalid parameter combinations before GPU operations."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/cuvs/neighbors/scann.hpp` around lines 178 - 196, The index constructor currently forwards params.n_coarse_clusters without validation; add input checks in index(raft::resources const& res, const index_params& params, IdxT n_rows, IdxT dim) to validate params.n_coarse_clusters (e.g., ensure it's >= 0 and, if >0, not greater than params.n_leaves or other incompatible combos) before performing any allocations or GPU operations; on invalid values use RAFT_EXPECTS with a clear message referencing n_coarse_clusters and n_leaves so the guard prevents entering training/allocation paths (update the RAFT_EXPECTS block alongside the existing pq_bits/pq_dim checks).
140-150:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPublic constructor change is source-breaking without deprecation path.
Adding required
n_coarse_clustersto this public constructor breaks existing callers that instantiatescann::indexdirectly with the previous signature. Please add a compatibility overload (or deprecation path + migration note) before removing the old form.As per coding guidelines, "for public C++ API headers ... Breaking changes require deprecation warnings and migration guide updates."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/cuvs/neighbors/scann.hpp` around lines 140 - 150, The new required parameter n_coarse_clusters makes the public constructor index(raft::resources const& res, cuvs::distance::DistanceType metric, uint32_t n_leaves, uint32_t n_coarse_clusters, ...) source-breaking; restore backwards compatibility by adding an overload for the previous signature (the constructor without n_coarse_clusters) that forwards to the new constructor with a sensible default for n_coarse_clusters, and mark that overload [[deprecated]] with a clear deprecation message pointing to the new parameter and migration steps; update the declaration of the existing constructor overloads in class scann::index (the index(...) constructor) so callers continue to build while receiving a deprecation warning and documentation pointer.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/cuvs/neighbors/scann.hpp`:
- Around line 250-260: Fix the typos and expand Doxygen on the public accessor
coarse_soar_labels(): change misspellings ("assighment" -> "assignment", "cosnt"
-> "const", and remove any stray tokens like "bg16"), and update the comment to
be full Doxygen describing the function purpose, return value
(raft::device_vector_view<uint32_t, IdxT>), and any side effects; also apply the
same spelling and Doxygen fixes to the adjacent public accessor comments that
return views (the other spilled-assignment accessor added nearby) so generated
docs are complete and consistent.
---
Outside diff comments:
In `@cpp/include/cuvs/neighbors/scann.hpp`:
- Around line 178-196: The index constructor currently forwards
params.n_coarse_clusters without validation; add input checks in
index(raft::resources const& res, const index_params& params, IdxT n_rows, IdxT
dim) to validate params.n_coarse_clusters (e.g., ensure it's >= 0 and, if >0,
not greater than params.n_leaves or other incompatible combos) before performing
any allocations or GPU operations; on invalid values use RAFT_EXPECTS with a
clear message referencing n_coarse_clusters and n_leaves so the guard prevents
entering training/allocation paths (update the RAFT_EXPECTS block alongside the
existing pq_bits/pq_dim checks).
- Around line 140-150: The new required parameter n_coarse_clusters makes the
public constructor index(raft::resources const& res,
cuvs::distance::DistanceType metric, uint32_t n_leaves, uint32_t
n_coarse_clusters, ...) source-breaking; restore backwards compatibility by
adding an overload for the previous signature (the constructor without
n_coarse_clusters) that forwards to the new constructor with a sensible default
for n_coarse_clusters, and mark that overload [[deprecated]] with a clear
deprecation message pointing to the new parameter and migration steps; update
the declaration of the existing constructor overloads in class scann::index (the
index(...) constructor) so callers continue to build while receiving a
deprecation warning and documentation pointer.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8aedfc39-6860-4bc8-9cc6-07fe89628892
📒 Files selected for processing (4)
cpp/include/cuvs/neighbors/scann.hppcpp/src/neighbors/scann/detail/scann_build.cuhcpp/src/neighbors/scann/detail/scann_serialize.cuhnotebooks/scann_protobuf_example.ipynb
🚧 Files skipped from review as they are similar to previous changes (3)
- cpp/src/neighbors/scann/detail/scann_serialize.cuh
- notebooks/scann_protobuf_example.ipynb
- cpp/src/neighbors/scann/detail/scann_build.cuh
Comment on lines
+250
to
+260
| /** spilled assighment of leaf centers to coarse centers by minimizing | ||
| * SOAR loss in two-level tree | ||
| */ | ||
| raft::device_vector_view<uint32_t, IdxT> coarse_soar_labels() noexcept | ||
| { | ||
| return coarse_soar_labels_.view(); | ||
| } | ||
|
|
||
| /** cosnt spilled assighment of leaf centers to coarse centers by minimizing | ||
| * SOAR loss in two-level tree | ||
| */ |
There was a problem hiding this comment.
Fix typos in public API docs to avoid degraded generated documentation.
There are multiple typos in newly added public comments (e.g., “assighment”, “cosnt”, “bg16”). Please correct them before merge.
As per coding guidelines, "All public API functions must include complete Doxygen documentation describing parameters, return values, and any side effects."
Also applies to: 318-319
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/cuvs/neighbors/scann.hpp` around lines 250 - 260, Fix the typos
and expand Doxygen on the public accessor coarse_soar_labels(): change
misspellings ("assighment" -> "assignment", "cosnt" -> "const", and remove any
stray tokens like "bg16"), and update the comment to be full Doxygen describing
the function purpose, return value (raft::device_vector_view<uint32_t, IdxT>),
and any side effects; also apply the same spelling and Doxygen fixes to the
adjacent public accessor comments that return views (the other
spilled-assignment accessor added nearby) so generated docs are complete and
consistent.
huuanhhuyn
requested changes
Jul 9, 2026
huuanhhuyn
left a comment
huuanhhuyn
left a comment
Contributor
There was a problem hiding this comment.
Thank you for the PR! I would have some minor doubts.
Should we / do we already have an example script or unit test to demonstrate how 2-level kmeans behaves?
rmaschal
force-pushed
the
scann-two-level
branch
from
July 14, 2026 20:28
0c8c0f7 to
4351585
Compare
rmaschal
force-pushed
the
scann-two-level
branch
2 times, most recently
from
July 14, 2026 20:37
e9dee86 to
65f335d
Compare
Contributor
Author
|
@huuanhhuyn thanks for the comments. I resolved the issues and added an additional examples for scann, using the the two-level funcionality. |
Contributor
|
/ok to test 1241a5c |
cjnolet
approved these changes
Jul 18, 2026
rmaschal
force-pushed
the
scann-two-level
branch
from
August 12, 2026 18:37
1241a5c to
d07c4aa
Compare
msarahan
approved these changes
Sep 8, 2026
Contributor
|
/ok to test 7b378d6 |
Contributor
|
/ok to test 7b378d6 |
@cjnolet, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Contributor
|
/ok to test f9ef48f |
Contributor
|
/ok to test 332b602 |
Contributor
Author
|
/ok to test 70140cb |
Contributor
|
/ok to test 1d9feb4 |
Contributor
|
/ok to test 9671660 |
Contributor
|
/ok to test 44eb92b |
@cjnolet, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Contributor
|
/ok to test 632ee13 |
Contributor
|
/merge |
bdice
added a commit
to bdice/cuvs
that referenced
this pull request
Oct 2, 2026
The two-level k-means tree training added in NVIDIA#1878 introduced a new rmm::device_async_resource_ref use after this branch was created.
bdice
added a commit
to bdice/cuvs
that referenced
this pull request
Oct 4, 2026
The two-level k-means tree training added in NVIDIA#1878 introduced a new rmm::device_async_resource_ref use after this branch was created.
copy-pr-bot Bot
pushed a commit
that referenced
this pull request
Oct 5, 2026
…pp-api Resolve conflicts with #1878 (two-level KMeans trees in ScaNN): - scann_build.cuh: keep the public cuvs::cluster::soar::detail::compute_soar_labels with upstream's centers_view rename. Also qualify the new coarse-level SOAR call #1878 added, which referenced the removed scann_soar.cuh helper. - examples/cpp/CMakeLists.txt: keep both SCANN_TWO_LEVEL_EXAMPLE and SOAR_EXAMPLE. Signed-off-by: Ranjit Rajan <ranjitr@nvidia.com>
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
7 participants
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.
This PR adds the option for a second level in the kmeans tree for the ScaNN index. The second level is built in a bottom-up fashion, e.g. second-level cluster centers are trained on leaf centers of the bottom level. AVQ and SOAR are also applied to second-level node centers.
One divergence from OSS ScaNN is that the second-level centers are trained on the normalized leaf centers, rather than directly on leaf centers. I've found this gives better recall than training directly on leaf centers, and often gives better recall compared to the OSS algo.
I've also added a notebook giving example code for producing missing protobuf artifacts for using cuVS built ScaNN indices with OSS ScaNN search functionality.