From 50d8b48f642ee34820eba9da25b83e31165a8910 Mon Sep 17 00:00:00 2001 From: Ranjit Rajan Date: Mon, 31 Aug 2026 17:20:02 -0700 Subject: [PATCH 1/2] Promote SOAR from a ScaNN internal to a public cluster API 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 --- cpp/CMakeLists.txt | 1 + cpp/include/cuvs/cluster/soar.hpp | 132 +++++++ .../detail/soar.cuh} | 77 +++- cpp/src/cluster/soar.cu | 54 +++ .../neighbors/scann/detail/scann_build.cuh | 19 +- cpp/tests/CMakeLists.txt | 1 + cpp/tests/cluster/soar.cu | 368 ++++++++++++++++++ fern/docs.yml | 2 + fern/pages/cpp_api/cpp-api-cluster-soar.md | 69 ++++ fern/pages/cpp_api/index.md | 1 + 10 files changed, 701 insertions(+), 23 deletions(-) create mode 100644 cpp/include/cuvs/cluster/soar.hpp rename cpp/src/{neighbors/scann/detail/scann_soar.cuh => cluster/detail/soar.cuh} (68%) create mode 100644 cpp/src/cluster/soar.cu create mode 100644 cpp/tests/cluster/soar.cu create mode 100644 fern/pages/cpp_api/cpp-api-cluster-soar.md diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 69ede8404e..1f9bc7a957 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1358,6 +1358,7 @@ if(NOT BUILD_CPU_ONLY) src/cluster/kmeans_transform_double.cu src/cluster/kmeans_transform_float.cu src/cluster/single_linkage_float.cu + src/cluster/soar.cu src/cluster/spectral.cu src/core/bitset.cu src/core/bloom_filter.cu diff --git a/cpp/include/cuvs/cluster/soar.hpp b/cpp/include/cuvs/cluster/soar.hpp new file mode 100644 index 0000000000..4c082e9ce4 --- /dev/null +++ b/cpp/include/cuvs/cluster/soar.hpp @@ -0,0 +1,132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +namespace CUVS_EXPORT cuvs { +namespace cluster { +namespace soar { + +/** + * @defgroup soar_params SOAR hyperparameters + * @{ + */ + +/** + * Simple object to specify hyper-parameters for SOAR assignment. + */ +struct params { + /** + * Weight of the projection of the secondary residual onto the primary residual in the SOAR + * loss. Larger values penalize secondary centroids whose residual is aligned with the primary + * residual, favoring complementary assignments. `0` reduces the loss to plain squared distance, + * which the primary centroid itself minimizes, so nothing is spilled. Default: 1.0. + */ + float lambda = 1.0f; +}; + +/** + * @} + */ + +/** + * @defgroup soar_predict SOAR assignment + * @{ + */ + +/** + * @brief Assign a secondary ("spilled") cluster to each row of the dataset. + * + * SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each vector, a second + * centroid that complements the primary assignment instead of merely being the next-closest + * one. It minimizes the loss of Theorem 3.1 of https://arxiv.org/abs/2404.00774: for a vector + * `x` with primary residual `r = x - centroids[labels[i]]`, + * + * `score(c) = ||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2` + * + * and `soar_labels[i]` is the centroid minimizing that score. Indexing a vector under both its + * primary and its secondary centroid improves recall for queries near a partition boundary. + * + * Only float32 data and uint32 labels are supported. + * + * The primary centroid is not excluded from the search, so `soar_labels[i] == labels[i]` is a + * possible (and meaningful) result: it says that no other centroid is worth spilling to, which + * is the common case for vectors in the interior of a cluster. Callers that treat SOAR as a + * strictly second posting list should test for this case and skip those rows. + * + * Scratch memory scales as `n_rows * n_clusters * 4` bytes because scores against all centroids + * are materialized at once and are not tiled. Process the dataset in row batches to bound the + * peak device memory usage. + * + * @code{.cpp} + * #include + * #include + * #include + * using namespace cuvs::cluster; + * ... + * raft::resources handle; + * cuvs::cluster::kmeans::balanced_params kmeans_params; + * int64_t n_features = 15, n_clusters = 100; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * + * // primary assignments, e.g. from balanced k-means + * kmeans::fit(handle, + * kmeans_params, + * dataset, + * centroids.view()); + * ... + * auto labels = raft::make_device_vector(handle, dataset.extent(0)); + * + * kmeans::predict(handle, + * kmeans_params, + * dataset, + * raft::make_const_mdspan(centroids.view()), + * labels.view()); + * ... + * // secondary assignments + * cuvs::cluster::soar::params soar_params; + * auto soar_labels = raft::make_device_vector(handle, dataset.extent(0)); + * + * soar::predict(handle, + * soar_params, + * dataset, + * raft::make_const_mdspan(centroids.view()), + * raft::make_const_mdspan(labels.view()), + * soar_labels.view()); + * // soar_labels now holds one secondary centroid id per row + * @endcode + * + * @param[in] handle The raft handle. + * @param[in] params Parameters for SOAR assignment. + * @param[in] dataset The dataset. The data must be in row-major format. + * [dim = n_rows x n_features] + * @param[in] centroids Cluster centroids. The data must be in row-major format. + * [dim = n_clusters x n_features] + * @param[in] labels Index of the primary cluster each row belongs to, as produced by + * k-means prediction. Every value must be in `[0, n_clusters)`. + * [len = n_rows] + * @param[out] soar_labels Index of the secondary cluster each row is spilled to. + * [len = n_rows] + */ +void predict(raft::resources const& handle, + const soar::params& params, + raft::device_matrix_view dataset, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + raft::device_vector_view soar_labels); + +/** + * @} + */ + +} // namespace soar +} // namespace cluster +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/neighbors/scann/detail/scann_soar.cuh b/cpp/src/cluster/detail/soar.cuh similarity index 68% rename from cpp/src/neighbors/scann/detail/scann_soar.cuh rename to cpp/src/cluster/detail/soar.cuh index 38bb7b0858..b039dd3e33 100644 --- a/cpp/src/neighbors/scann/detail/scann_soar.cuh +++ b/cpp/src/cluster/detail/soar.cuh @@ -1,13 +1,16 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +#pragma once + #include #include #include #include #include +#include #include #include #include @@ -17,27 +20,72 @@ #include #include +namespace cuvs::cluster::soar::detail { + +/** + * @brief Subtract cluster center coordinates from each dataset vector. + * + * residual[i, k] = dataset[i ,k] - centers[l, k], + * where l = labels[i], the cluster label corresponding to vector i. + * + * An identical copy lives in `neighbors/scann/detail/scann_quantize.cuh`, where it is also used + * outside the SOAR path to build the PQ trainset residuals. + * + * @tparam T + * @tparam LabelT + * @param res raft resources + * @param dataset dataset vectors, size [n_rows, dim] + * @param centers cluster center coordinates, size [n_clusters, dim] + * @param labels cluster labels, size [n_rows] + * @return device matrix with the residuals, size [n_rows, dim] + */ +template +auto compute_residuals(raft::resources const& res, + raft::device_matrix_view dataset, + raft::device_matrix_view centers, + raft::device_vector_view labels) + -> raft::device_matrix +{ + auto dim = dataset.extent(1); + auto residuals = raft::make_device_matrix(res, labels.extent(0), dim); + + raft::linalg::map_offset( + res, residuals.view(), [dataset, centers, labels, dim] __device__(size_t i) { + int row_idx = i / dim; + int el_idx = i % dim; + return dataset(row_idx, el_idx) - centers(labels(row_idx), el_idx); + }); + + return residuals; +} + /** * @brief Compute SOAR labels for each dataset vector * * Compute a second, spilled cluster for each dataset vector by minimizing * the loss function in Theorem 3.1 of https://arxiv.org/abs/2404.00774 * + * Residuals are an input (`r = x - centers[labels[i]]`) rather than derived here, so a + * caller that already has them can avoid a second pass over the dataset. + * + * The scratch score matrix is [n_rows, n_clusters] floats and is not tiled, so callers are + * responsible for batching rows. + * * @tparam T - * @tparam LavelT - * @param res raft resources + * @tparam LabelT + * @param dev_resources raft resources * @param dataset the dataset, size [n_rows, dim] * @param residuals the residual vectors r, size [n_rows, dim] * @param centers the cluster centers, size [n_clusters, dim] * @param labels the cluster assignments, size [n_rows] - * @param soar_labels the computed soar labels + * @param soar_labels the computed soar labels, size [n_rows] * @param lambda the weight for the projection of a residual r' onto r in the SOAR loss */ template void compute_soar_labels(raft::resources const& dev_resources, raft::device_matrix_view dataset, raft::device_matrix_view residuals, - raft::device_matrix_view centers, + raft::device_matrix_view centers, raft::device_vector_view labels, raft::device_vector_view soar_labels, float lambda) @@ -47,7 +95,6 @@ void compute_soar_labels(raft::resources const& dev_resources, // compute SOAR metric for each center auto soar_scores = raft::make_device_matrix(dev_resources, dataset.extent(0), centers.extent(0)); - auto n_centers = centers.extent(0); auto residuals_norm = raft::make_device_matrix( dev_resources, residuals.extent(0), residuals.extent(1)); @@ -90,15 +137,15 @@ void compute_soar_labels(raft::resources const& dev_resources, auto centers_transpose = raft::make_device_matrix(dev_resources, centers.extent(1), centers.extent(0)); - raft::linalg::reduce(dev_resources, - raft::make_const_mdspan(centers), - centers_norm.view(), - 0.0f, - false, - raft::sq_op(), - raft::add_op()); + raft::linalg::reduce( + dev_resources, centers, centers_norm.view(), 0.0f, false, raft::sq_op(), raft::add_op()); - raft::linalg::transpose(dev_resources, centers, centers_transpose.view()); + // raft::linalg::transpose requires input and output views of the same type; it does not + // write to the input. + auto nc_centers = raft::make_device_matrix_view( + const_cast(centers.data_handle()), centers.extent(0), centers.extent(1)); + + raft::linalg::transpose(dev_resources, nc_centers, centers_transpose.view()); raft::linalg::gemm( dev_resources, residuals_norm.view(), centers_transpose.view(), soar_scores.view()); @@ -146,3 +193,5 @@ void compute_soar_labels(raft::resources const& dev_resources, raft::matrix::argmin(dev_resources, raft::make_const_mdspan(soar_scores.view()), soar_labels); } + +} // namespace cuvs::cluster::soar::detail diff --git a/cpp/src/cluster/soar.cu b/cpp/src/cluster/soar.cu new file mode 100644 index 0000000000..804755003f --- /dev/null +++ b/cpp/src/cluster/soar.cu @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "detail/soar.cuh" + +#include + +#include + +namespace cuvs::cluster::soar { + +void predict(raft::resources const& handle, + const soar::params& params, + raft::device_matrix_view dataset, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + raft::device_vector_view soar_labels) +{ + int64_t n_rows = dataset.extent(0); + int64_t dim = dataset.extent(1); + int64_t n_clusters = centroids.extent(0); + + RAFT_EXPECTS(centroids.extent(1) == dim, + "Number of features in the dataset (%zd) and in the centroids (%zd) must match.", + dim, + centroids.extent(1)); + RAFT_EXPECTS(n_clusters > 0, "The number of centroids must be positive."); + RAFT_EXPECTS(dim > 0, "The number of features must be positive."); + RAFT_EXPECTS(labels.extent(0) == n_rows, + "The number of labels (%zd) must match the number of rows in the dataset (%zd).", + labels.extent(0), + n_rows); + RAFT_EXPECTS( + soar_labels.extent(0) == n_rows, + "The number of soar labels (%zd) must match the number of rows in the dataset (%zd).", + soar_labels.extent(0), + n_rows); + + if (n_rows == 0) { return; } + + auto residuals = detail::compute_residuals(handle, dataset, centroids, labels); + + detail::compute_soar_labels(handle, + dataset, + raft::make_const_mdspan(residuals.view()), + centroids, + labels, + soar_labels, + params.lambda); +} + +} // namespace cuvs::cluster::soar diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index c01e50bc83..5f40f2c510 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -24,10 +24,10 @@ #include #include +#include "../../../cluster/detail/soar.cuh" #include "scann_avq.cuh" #include "scann_common.cuh" #include "scann_quantize.cuh" -#include "scann_soar.cuh" namespace cuvs::neighbors::experimental::scann::detail { using namespace cuvs::spatial::knn::detail; // NOLINT @@ -197,13 +197,14 @@ index build( // Compute SOAR labels. // We compute SOAR labels in this loop to eliminate one HtoD copy of the full dataset. - compute_soar_labels(res, - batch_view, - raft::make_const_mdspan(avq_residuals.view()), - centroids_view, - batch_labels_view, - batch_soar_labels_view, - params.soar_lambda); + cuvs::cluster::soar::detail::compute_soar_labels( + res, + batch_view, + raft::make_const_mdspan(avq_residuals.view()), + raft::make_const_mdspan(centroids_view), + batch_labels_view, + batch_soar_labels_view, + params.soar_lambda); // Compute and quantize residuals using the public PQ API int64_t codes_dim = cuvs::preprocessing::quantize::pq::get_quantized_dim(pq_build_params); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 744bc6a7a2..19a1a2b14b 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -131,6 +131,7 @@ ConfigureTest( cluster/kmeans_predict_batching.cu cluster/linkage.cu cluster/connect_knn.cu + cluster/soar.cu cluster/spectral.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/cluster/soar.cu b/cpp/tests/cluster/soar.cu new file mode 100644 index 0000000000..2a5598e32a --- /dev/null +++ b/cpp/tests/cluster/soar.cu @@ -0,0 +1,368 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../src/cluster/detail/soar.cuh" +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::cluster::soar { + +struct SoarInputs { + int64_t n_rows; + int64_t dim; + int64_t n_clusters; + float lambda; +}; + +::std::ostream& operator<<(::std::ostream& os, const SoarInputs& p) +{ + os << "{ " << p.n_rows << ", " << p.dim << ", " << p.n_clusters << ", " << p.lambda << '}'; + return os; +} + +namespace { + +/** Uniform random matrix in [-1, 1], generated on the host so the tests are reproducible. */ +auto random_matrix(int64_t n_rows, int64_t dim, uint64_t seed) -> std::vector +{ + std::mt19937_64 rng(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector data(n_rows * dim); + std::generate(data.begin(), data.end(), [&]() { return dist(rng); }); + return data; +} + +/** Index of the closest centroid in L2, i.e. what k-means prediction would produce. */ +auto nearest_centroids(const std::vector& dataset, + const std::vector& centroids, + int64_t n_rows, + int64_t dim, + int64_t n_clusters) -> std::vector +{ + std::vector labels(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + double best_distance = std::numeric_limits::max(); + for (int64_t c = 0; c < n_clusters; c++) { + double distance = 0.0; + for (int64_t k = 0; k < dim; k++) { + double diff = static_cast(dataset[i * dim + k]) - centroids[c * dim + k]; + distance += diff * diff; + } + if (distance < best_distance) { + best_distance = distance; + labels[i] = static_cast(c); + } + } + } + return labels; +} + +auto residuals_host(const std::vector& dataset, + const std::vector& centroids, + const std::vector& labels, + int64_t n_rows, + int64_t dim) -> std::vector +{ + std::vector residuals(n_rows * dim); + for (int64_t i = 0; i < n_rows; i++) { + for (int64_t k = 0; k < dim; k++) { + residuals[i * dim + k] = dataset[i * dim + k] - centroids[labels[i] * dim + k]; + } + } + return residuals; +} + +/** + * `||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2`, the loss that the device implementation + * minimizes over all centroids, up to a per-row constant that does not move the argmin. + */ +auto soar_score( + const float* x, const float* residual, const float* centroid, int64_t dim, float lambda) -> double +{ + double residual_norm = 0.0; + for (int64_t k = 0; k < dim; k++) { + residual_norm += static_cast(residual[k]) * residual[k]; + } + residual_norm = std::sqrt(residual_norm); + + double squared_distance = 0.0; + double projection = 0.0; + for (int64_t k = 0; k < dim; k++) { + double diff = static_cast(x[k]) - centroid[k]; + squared_distance += diff * diff; + projection += diff * (residual[k] / residual_norm); + } + + return squared_distance + static_cast(lambda) * projection * projection; +} + +/** The best achievable loss per row, from an exhaustive host search. */ +auto reference_scores(const std::vector& dataset, + const std::vector& centroids, + const std::vector& residuals, + int64_t n_rows, + int64_t dim, + int64_t n_clusters, + float lambda) -> std::vector +{ + std::vector scores(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + double best_score = std::numeric_limits::max(); + for (int64_t c = 0; c < n_clusters; c++) { + best_score = std::min( + best_score, + soar_score(&dataset[i * dim], &residuals[i * dim], ¢roids[c * dim], dim, lambda)); + } + scores[i] = best_score; + } + return scores; +} + +} // namespace + +class SoarTest : public ::testing::TestWithParam { + public: + SoarTest() + : params_(GetParam()), + dataset_(raft::make_device_matrix(handle_, params_.n_rows, params_.dim)), + centroids_( + raft::make_device_matrix(handle_, params_.n_clusters, params_.dim)), + labels_(raft::make_device_vector(handle_, params_.n_rows)), + soar_labels_(raft::make_device_vector(handle_, params_.n_rows)) + { + } + + protected: + void SetUp() override + { + h_dataset_ = random_matrix(params_.n_rows, params_.dim, 1234ULL); + h_centroids_ = random_matrix(params_.n_clusters, params_.dim, 5678ULL); + h_labels_ = + nearest_centroids(h_dataset_, h_centroids_, params_.n_rows, params_.dim, params_.n_clusters); + + auto stream = raft::resource::get_cuda_stream(handle_); + raft::update_device(dataset_.data_handle(), h_dataset_.data(), h_dataset_.size(), stream); + raft::update_device(centroids_.data_handle(), h_centroids_.data(), h_centroids_.size(), stream); + raft::update_device(labels_.data_handle(), h_labels_.data(), h_labels_.size(), stream); + raft::resource::sync_stream(handle_); + } + + /** Run the public API and copy the resulting labels back to the host. */ + auto run_predict() -> std::vector + { + cuvs::cluster::soar::params soar_params; + soar_params.lambda = params_.lambda; + + cuvs::cluster::soar::predict(handle_, + soar_params, + raft::make_const_mdspan(dataset_.view()), + raft::make_const_mdspan(centroids_.view()), + raft::make_const_mdspan(labels_.view()), + soar_labels_.view()); + + return to_host(raft::make_const_mdspan(soar_labels_.view())); + } + + auto to_host(raft::device_vector_view labels) -> std::vector + { + std::vector h_labels(labels.extent(0)); + raft::update_host(h_labels.data(), + labels.data_handle(), + labels.extent(0), + raft::resource::get_cuda_stream(handle_)); + raft::resource::sync_stream(handle_); + return h_labels; + } + + raft::resources handle_; + SoarInputs params_; + + std::vector h_dataset_; + std::vector h_centroids_; + std::vector h_labels_; + + raft::device_matrix dataset_; + raft::device_matrix centroids_; + raft::device_vector labels_; + raft::device_vector soar_labels_; +}; + +/** + * Every label must be a valid centroid id achieving the same loss as an exhaustive host search. + * Comparing losses rather than ids keeps the test from being fragile when two centroids are + * nearly tied. + */ +TEST_P(SoarTest, MatchesHostReference) +{ + auto soar_labels = run_predict(); + + auto h_residuals = + residuals_host(h_dataset_, h_centroids_, h_labels_, params_.n_rows, params_.dim); + auto best_scores = reference_scores(h_dataset_, + h_centroids_, + h_residuals, + params_.n_rows, + params_.dim, + params_.n_clusters, + params_.lambda); + + for (int64_t i = 0; i < params_.n_rows; i++) { + ASSERT_LT(soar_labels[i], static_cast(params_.n_clusters)) + << "label out of range at row " << i; + + double score = soar_score(&h_dataset_[i * params_.dim], + &h_residuals[i * params_.dim], + &h_centroids_[soar_labels[i] * params_.dim], + params_.dim, + params_.lambda); + ASSERT_NEAR(score, best_scores[i], 1e-4 * (1.0 + std::abs(best_scores[i]))) + << "row " << i << " picked centroid " << soar_labels[i]; + } +} + +/** + * The residuals feeding the SOAR loss must match a plain host `x - c[label]`. Both the public + * `predict` and the ScaNN builder depend on this, and `MatchesHostReference` only detects + * residual errors large enough to move an argmin. + */ +TEST_P(SoarTest, ComputeResidualsMatchesHost) +{ + auto h_residuals = + residuals_host(h_dataset_, h_centroids_, h_labels_, params_.n_rows, params_.dim); + + auto residuals = + detail::compute_residuals(handle_, + raft::make_const_mdspan(dataset_.view()), + raft::make_const_mdspan(centroids_.view()), + raft::make_const_mdspan(labels_.view())); + + ASSERT_TRUE(cuvs::devArrMatchHost(h_residuals.data(), + residuals.data_handle(), + h_residuals.size(), + cuvs::CompareApprox(1e-6f), + raft::resource::get_cuda_stream(handle_))); +} + +const std::vector inputs = {{1000, 8, 16, 1.0f}, + {1000, 8, 16, 0.0f}, + {1000, 8, 16, 4.0f}, + {512, 32, 64, 1.5f}, + {17, 3, 2, 1.0f}}; + +INSTANTIATE_TEST_CASE_P(SoarTests, SoarTest, ::testing::ValuesIn(inputs)); + +/** + * A hand-checked case with well-separated centroids, covering both outcomes: a row in the + * interior of its cluster keeps its primary centroid, because no other centroid is close enough + * to be worth spilling to, while a row near a boundary spills to the neighboring cluster. + */ +TEST(SoarTestSmall, SeparatedClusters) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t n_rows = 4, dim = 2, n_clusters = 3; + + std::vector h_centroids{0.0f, 0.0f, 100.0f, 0.0f, 0.0f, 100.0f}; + std::vector h_dataset{1.0f, 0.0f, 48.0f, 20.0f, 20.0f, 48.0f, 99.0f, 0.0f}; + std::vector h_labels{0, 0, 0, 1}; + + // Rows 0 and 3 sit next to their own centroid and keep it. Rows 1 and 2 sit between two + // centroids, so the second-closest one wins with a ~13% margin in the loss. + std::vector expected{0, 1, 2, 1}; + + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + auto centroids = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + auto soar_labels = raft::make_device_vector(handle, n_rows); + + raft::update_device(dataset.data_handle(), h_dataset.data(), h_dataset.size(), stream); + raft::update_device(centroids.data_handle(), h_centroids.data(), h_centroids.size(), stream); + raft::update_device(labels.data_handle(), h_labels.data(), h_labels.size(), stream); + + cuvs::cluster::soar::params params; + cuvs::cluster::soar::predict(handle, + params, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(centroids.view()), + raft::make_const_mdspan(labels.view()), + soar_labels.view()); + + std::vector result(n_rows); + raft::update_host(result.data(), soar_labels.data_handle(), n_rows, stream); + raft::resource::sync_stream(handle); + + ASSERT_EQ(expected, result); +} + +TEST(SoarTestErrors, RejectsMismatchedShapes) +{ + raft::resources handle; + + constexpr int64_t n_rows = 32, dim = 4, n_clusters = 8; + + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + auto centroids = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + auto soar_labels = raft::make_device_vector(handle, n_rows); + + cuvs::cluster::soar::params params; + auto dataset_view = raft::make_const_mdspan(dataset.view()); + auto centroids_view = raft::make_const_mdspan(centroids.view()); + auto labels_view = raft::make_const_mdspan(labels.view()); + + // centroid dimension differs from the dataset dimension + auto narrow_centroids = raft::make_device_matrix(handle, n_clusters, dim - 1); + ASSERT_THROW(cuvs::cluster::soar::predict(handle, + params, + dataset_view, + raft::make_const_mdspan(narrow_centroids.view()), + labels_view, + soar_labels.view()), + raft::logic_error); + + // no centroids to choose from + auto no_centroids = raft::make_device_matrix(handle, 0, dim); + ASSERT_THROW(cuvs::cluster::soar::predict(handle, + params, + dataset_view, + raft::make_const_mdspan(no_centroids.view()), + labels_view, + soar_labels.view()), + raft::logic_error); + + // one primary label per row is required + auto short_labels = raft::make_device_vector(handle, n_rows - 1); + ASSERT_THROW(cuvs::cluster::soar::predict(handle, + params, + dataset_view, + centroids_view, + raft::make_const_mdspan(short_labels.view()), + soar_labels.view()), + raft::logic_error); + + // one output slot per row is required + auto short_output = raft::make_device_vector(handle, n_rows - 1); + ASSERT_THROW(cuvs::cluster::soar::predict( + handle, params, dataset_view, centroids_view, labels_view, short_output.view()), + raft::logic_error); +} + +} // namespace cuvs::cluster::soar diff --git a/fern/docs.yml b/fern/docs.yml index 2aafaea69e..1f846f0d43 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -327,6 +327,8 @@ navigation: path: "./pages/cpp_api/cpp-api-cluster-gmm.md" - page: "Cluster Kmeans" path: "./pages/cpp_api/cpp-api-cluster-kmeans.md" + - page: "Cluster Soar" + path: "./pages/cpp_api/cpp-api-cluster-soar.md" - page: "Cluster Spectral" path: "./pages/cpp_api/cpp-api-cluster-spectral.md" - section: "Common Types" diff --git a/fern/pages/cpp_api/cpp-api-cluster-soar.md b/fern/pages/cpp_api/cpp-api-cluster-soar.md new file mode 100644 index 0000000000..a5712a0c08 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-cluster-soar.md @@ -0,0 +1,69 @@ +--- +slug: api-reference/cpp-api-cluster-soar +--- + +# Soar + +_Source header: `cuvs/cluster/soar.hpp`_ + +## SOAR hyperparameters + + +### cluster::soar::params + +Simple object to specify hyper-parameters for SOAR assignment. + +```cpp +struct params { + float lambda = 1.0f; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `lambda` | `float` | Weight of the projection of the secondary residual onto the primary residual in the SOAR loss. Larger values penalize secondary centroids whose residual is aligned with the primary residual, favoring complementary assignments. `0` reduces the loss to plain squared distance, which the primary centroid itself minimizes, so nothing is spilled.
Default: `1.0`. | + +## SOAR assignment + + +### cluster::soar::predict + +Assign a secondary ("spilled") cluster to each row of the dataset. + +```cpp +void predict(raft::resources const& handle, +const soar::params& params, +raft::device_matrix_view dataset, +raft::device_matrix_view centroids, +raft::device_vector_view labels, +raft::device_vector_view soar_labels); +``` + +SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each vector, a second centroid that complements the primary assignment instead of merely being the next-closest one. It minimizes the loss of Theorem 3.1 of https://arxiv.org/abs/2404.00774: for a vector `x` with primary residual `r = x - centroids[labels[i]]`, + +`score(c) = ||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2` + +and `soar_labels[i]` is the centroid minimizing that score. Indexing a vector under both its primary and its secondary centroid improves recall for queries near a partition boundary. + +Only float32 data and uint32 labels are supported. + +The primary centroid is not excluded from the search, so `soar_labels[i] == labels[i]` is a possible (and meaningful) result: it says that no other centroid is worth spilling to, which is the common case for vectors in the interior of a cluster. Callers that treat SOAR as a strictly second posting list should test for this case and skip those rows. + +Scratch memory scales as `n_rows * n_clusters * 4` bytes because scores against all centroids are materialized at once and are not tiled. Process the dataset in row batches to bound the peak device memory usage. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | The raft handle. | +| `params` | in | [`const soar::params&`](/api-reference/cpp-api-cluster-soar#cluster-soar-params) | Parameters for SOAR assignment. | +| `dataset` | in | `raft::device_matrix_view` | The dataset. The data must be in row-major format. [dim = n_rows x n_features] | +| `centroids` | in | `raft::device_matrix_view` | Cluster centroids. The data must be in row-major format. [dim = n_clusters x n_features] | +| `labels` | in | `raft::device_vector_view` | Index of the primary cluster each row belongs to, as produced by k-means prediction. Every value must be in `[0, n_clusters)`. [len = n_rows] | +| `soar_labels` | out | `raft::device_vector_view` | Index of the secondary cluster each row is spilled to. [len = n_rows] | + +**Returns** + +`void` diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index d5904b5db4..bd3364d987 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -7,6 +7,7 @@ These pages are generated from the documented public headers in the cuVS source - [Agglomerative](/api-reference/cpp-api-cluster-agglomerative) - [Gmm](/api-reference/cpp-api-cluster-gmm) - [K-Means](/api-reference/cpp-api-cluster-kmeans) +- [Soar](/api-reference/cpp-api-cluster-soar) - [Spectral](/api-reference/cpp-api-cluster-spectral) ## Common From 91b54cbaf7103591f4d973a3b76a065c51dd0af4 Mon Sep 17 00:00:00 2001 From: Ranjit Rajan Date: Tue, 1 Sep 2026 10:45:28 -0700 Subject: [PATCH 2/2] Add a standalone C++ example for the SOAR clustering API `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 --- examples/cpp/CMakeLists.txt | 2 + examples/cpp/src/soar_example.cu | 102 +++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+) create mode 100644 examples/cpp/src/soar_example.cu diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index f4dffad3a9..09e64c87b9 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -51,6 +51,7 @@ add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) add_executable(IVF_PQ_EXAMPLE src/ivf_pq_example.cu) add_executable(VAMANA_EXAMPLE src/vamana_example.cu) add_executable(SCANN_EXAMPLE src/scann_example.cu) +add_executable(SOAR_EXAMPLE src/soar_example.cu) # `$` is a generator expression that ensures that targets are # installed in a conda environment, if one exists @@ -80,3 +81,4 @@ target_link_libraries(IVF_PQ_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(VAMANA_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(SCANN_EXAMPLE PRIVATE cuvs::cuvs $) +target_link_libraries(SOAR_EXAMPLE PRIVATE cuvs::cuvs $) diff --git a/examples/cpp/src/soar_example.cu b/examples/cpp/src/soar_example.cu new file mode 100644 index 0000000000..8821a164de --- /dev/null +++ b/examples/cpp/src/soar_example.cu @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +/** Number of rows whose secondary cluster differs from their primary one. */ +int64_t count_spilled(raft::device_resources const& dev_resources, + raft::device_vector_view labels, + raft::device_vector_view soar_labels) +{ + auto h_labels = raft::make_host_vector(labels.extent(0)); + auto h_soar_labels = raft::make_host_vector(soar_labels.extent(0)); + auto stream = raft::resource::get_cuda_stream(dev_resources); + + raft::copy(h_labels.data_handle(), labels.data_handle(), labels.size(), stream); + raft::copy(h_soar_labels.data_handle(), soar_labels.data_handle(), soar_labels.size(), stream); + raft::resource::sync_stream(dev_resources, stream); + + int64_t n_spilled = 0; + for (int64_t i = 0; i < labels.extent(0); ++i) { + if (h_soar_labels(i) != h_labels(i)) { ++n_spilled; } + } + return n_spilled; +} + +void soar_predict_example(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view centroids, + raft::device_vector_view labels) +{ + // Default lambda = 1. Larger values penalize secondary centroids whose residual is aligned with + // the primary residual, favoring more complementary assignments. + cuvs::cluster::soar::params params; + + auto soar_labels = raft::make_device_vector(dev_resources, dataset.extent(0)); + + cuvs::cluster::soar::predict( + dev_resources, params, dataset, centroids, labels, soar_labels.view()); + + // A row keeps its primary label when no other centroid is worth spilling to. + auto n_spilled = + count_spilled(dev_resources, labels, raft::make_const_mdspan(soar_labels.view())); + + std::cout << "Spilled " << n_spilled << " of " << dataset.extent(0) + << " rows to a secondary cluster" << std::endl; +} + +int main() +{ + raft::device_resources dev_resources; + + // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool. + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), + 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(pool_mr); + + int64_t n_samples = 10000; + int64_t n_dim = 64; + int64_t n_clusters = 100; + + // Far fewer blobs than k-means clusters: the 100 learned partitions subdivide the 10 dense + // regions, creating internal partition boundaries. blob_labels is required by make_blobs but + // unused. + int64_t n_blobs = 10; + auto dataset = raft::make_device_matrix(dev_resources, n_samples, n_dim); + auto blob_labels = raft::make_device_vector(dev_resources, n_samples); + raft::random::make_blobs(dev_resources, dataset.view(), blob_labels.view(), n_blobs); + + auto dataset_view = raft::make_const_mdspan(dataset.view()); + + // SOAR needs centroids and a primary label per row, so k-means runs first. + cuvs::cluster::kmeans::balanced_params kmeans_params; + auto centroids = raft::make_device_matrix(dev_resources, n_clusters, n_dim); + auto labels = raft::make_device_vector(dev_resources, n_samples); + + cuvs::cluster::kmeans::fit(dev_resources, kmeans_params, dataset_view, centroids.view()); + cuvs::cluster::kmeans::predict(dev_resources, + kmeans_params, + dataset_view, + raft::make_const_mdspan(centroids.view()), + labels.view()); + + soar_predict_example(dev_resources, + dataset_view, + raft::make_const_mdspan(centroids.view()), + raft::make_const_mdspan(labels.view())); +}