Repository navigation
Conversation
cjnolet
added
improvement
SOAR (Spilling with Orthogonality-Amplified Residuals, https://arxiv.org/abs/2404.00774) gives each vector a second centroid chosen to complement its primary assignment rather than to be merely the next closest. Indexing a vector under both partitions improves recall for queries near a partition boundary. The implementation lived in `neighbors/scann/detail/scann_soar.cuh` and was reachable only by building a ScaNN index, even though the algorithm needs nothing beyond centroids and primary k-means labels. Promotes it to a cluster-level API. `cuvs::cluster::soar::predict` takes a dataset, centroids, and primary labels, and writes one secondary label per row. `soar::params` exposes the `lambda` weight controlling how strongly a candidate centroid is penalized for having a residual aligned with the primary one. Moves `scann_soar.cuh` to `cluster/detail/soar.cuh` and points the ScaNN builder at the relocated entry point so there is a single implementation. `compute_soar_labels` now takes its centroids as a const view. The detail header also gains `compute_residuals`, which the public API needs to derive residuals from labels; the ScaNN builder already holds residuals for quantization and keeps supplying its own, so it does not pay for a second pass over the dataset. Adds `cpp/tests/cluster/soar.cu` to `CLUSTER_TEST`, covering assignments against an exhaustive host search, the residual computation against a host reference, a hand-checked separated-cluster case, and the shape-validation errors. Adds a C++ API documentation page. Signed-off-by: Ranjit Rajan <ranjitr@nvidia.com>
`SOAR_EXAMPLE` shows the call sequence a caller needs: balanced k-means for the centroids and primary labels, then `soar::predict` to fill one secondary label per row. Signed-off-by: Ranjit Rajan <ranjitr@nvidia.com>
cuvsSoarPredict exposes soar::predict through the C API, and cuvs.cluster.soar.predict wraps that for Python. Both accept the int32 labels written by cuvsKMeansPredict as well as uint32. The C and Python tests score the returned labels against an independent host evaluation of the SOAR loss, check that the int32 and uint32 paths agree, and reject mismatched dtypes and shapes. The Fern C and Python reference pages are generated from the new sources. Signed-off-by: Ranjit Rajan <ranjitr@nvidia.com>
ronjer30
force-pushed
the
feat/cluster-soar-cpp-api
branch
from
September 28, 2026 22:13
91b54cb to
0d3e683
Compare
raft::resource::get_cuda_stream now returns cuda::stream_ref, which does not convert implicitly to the cudaStream_t that cuvs::devArrMatchHost takes, so the conda-cpp-build and devcontainer CI jobs failed to compile cpp/tests/cluster/soar.cu. Pass the underlying stream with .get() Signed-off-by: Ranjit Rajan <ranjitr@nvidia.com>
Contributor
|
@ronjer30 looks like this PR accrued some merge conflicts. Can you fix? I'd like to get this reviewed. |
…pp-api Resolve conflicts with NVIDIA#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 NVIDIA#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>
rmm::cuda_stream_view was removed in rapidsai/rmm#2552, and NVIDIA#2710 removed the remaining uses from cuVS. SoarFixture still took one, so SOAR_C_TEST would fail to build against newer RMM nightlies. Take cuda::stream_ref instead, as the other tests do. Signed-off-by: Ranjit Rajan <ranjitr@nvidia.com>
Author
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
2 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.
Description
Testing
SOAR_EXAMPLE.SCANN_EXAMPLEindex build time on NVIDIA GB10, 10 runs after warm-up: 412.3 ± 4.0 ms before and 410.8 ± 6.0 ms after; no measurable regression.