[WIP] Expose SOAR as a public cluster API

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

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 <cuvs/core/export.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/resources.hpp>
+
+#include <cstdint>
+
+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 <raft/core/resources.hpp>
+ *   #include <cuvs/cluster/kmeans.hpp>
+ *   #include <cuvs/cluster/soar.hpp>
+ *   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<float, int64_t>(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<uint32_t, int64_t>(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<uint32_t, int64_t>(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<const float, int64_t> dataset,
+             raft::device_matrix_view<const float, int64_t> centroids,
+             raft::device_vector_view<const uint32_t, int64_t> labels,
+             raft::device_vector_view<uint32_t, int64_t> soar_labels);
+
+/**
+ * @}
+ */
+
+}  // namespace soar
+}  // namespace cluster
+}  // namespace CUVS_EXPORT cuvs

cpp/src/neighbors/scann/detail/scann_soar.cuh → 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 <raft/core/device_mdarray.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/error.hpp>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/cublas_handle.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/gemm.cuh>
 #include <raft/linalg/map.cuh>
@@ -17,27 +20,72 @@
 #include <raft/linalg/transpose.cuh>
 #include <raft/matrix/argmin.cuh>
 
+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 <typename T, typename LabelT>
+auto compute_residuals(raft::resources const& res,
+                       raft::device_matrix_view<const T, int64_t> dataset,
+                       raft::device_matrix_view<const T, int64_t> centers,
+                       raft::device_vector_view<const LabelT, int64_t> labels)
+  -> raft::device_matrix<T, int64_t, raft::row_major>
+{
+  auto dim       = dataset.extent(1);
+  auto residuals = raft::make_device_matrix<T, int64_t>(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 <typename T, typename LabelT>
 void compute_soar_labels(raft::resources const& dev_resources,
                          raft::device_matrix_view<const T, int64_t> dataset,
                          raft::device_matrix_view<const T, int64_t> residuals,
-                         raft::device_matrix_view<T, int64_t> centers,
+                         raft::device_matrix_view<const T, int64_t> centers,
                          raft::device_vector_view<const LabelT, int64_t> labels,
                          raft::device_vector_view<LabelT, int64_t> 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<float, int64_t>(dev_resources, dataset.extent(0), centers.extent(0));
-  auto n_centers = centers.extent(0);
 
   auto residuals_norm = raft::make_device_matrix<float, int64_t>(
     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<T, int64_t>(dev_resources, centers.extent(1), centers.extent(0));
 
-  raft::linalg::reduce<raft::Apply::ALONG_ROWS>(dev_resources,
-                                                raft::make_const_mdspan(centers),
-                                                centers_norm.view(),
-                                                0.0f,
-                                                false,
-                                                raft::sq_op(),
-                                                raft::add_op());
+  raft::linalg::reduce<raft::Apply::ALONG_ROWS>(
+    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<T, int64_t>(
+    const_cast<T*>(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

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 <cuvs/cluster/soar.hpp>
+
+#include <raft/core/error.hpp>
+
+namespace cuvs::cluster::soar {
+
+void predict(raft::resources const& handle,
+             const soar::params& params,
+             raft::device_matrix_view<const float, int64_t> dataset,
+             raft::device_matrix_view<const float, int64_t> centroids,
+             raft::device_vector_view<const uint32_t, int64_t> labels,
+             raft::device_vector_view<uint32_t, int64_t> 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<float, uint32_t>(handle, dataset, centroids, labels);
+
+  detail::compute_soar_labels<float, uint32_t>(handle,
+                                               dataset,
+                                               raft::make_const_mdspan(residuals.view()),
+                                               centroids,
+                                               labels,
+                                               soar_labels,
+                                               params.lambda);
+}
+
+}  // namespace cuvs::cluster::soar

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 <raft/matrix/slice.cuh>
 #include <raft/random/rng.cuh>
 
+#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<T, IdxT> build(
 
     // Compute SOAR labels.
     // We compute SOAR labels in this loop to eliminate one HtoD copy of the full dataset.
-    compute_soar_labels<T, uint32_t>(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<T, uint32_t>(
+      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);

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