[KMeans] Configure Workspace for Large Batch Sizes

c/include/cuvs/cluster/kmeans.h

@@ -106,6 +106,12 @@ struct cuvsKMeansParams {
   /**
    * Number of samples to process per GPU batch for the batched (host-data) API.
    * When set to 0, defaults to n_samples (process all at once).
+   *
+   * Host rows are staged through a device buffer of size
+   * `device_buffer_samples * n_features`. If the resources handle has a CUDA
+   * stream pool with at least one stream, prefetch is enabled; that
+   * doubles the batch memory footprint (and the weight-staging footprint
+   * when sample weights are provided).
    */
   int64_t device_buffer_samples;
 

cpp/include/cuvs/cluster/kmeans.hpp

@@ -150,7 +150,14 @@ struct params : base_params {
    *
    * In multi-GPU mode this is a per-rank batch size: each rank processes up
    * to this many local samples per batch, clamped to that rank's local sample
-   * count. This is is ignored by device-data overloads.
+   * count. This is ignored by device-data overloads.
+   *
+   * Host rows are staged through a device buffer of size
+   * `device_buffer_samples * n_features`. If the resources handle has a CUDA
+   * stream pool with at least one stream, prefetch is enabled; that
+   * doubles the batch memory footprint (and the weight-staging footprint
+   * when sample weights are provided).
+   *
    * Default: 0 (process all data at once).
    */
   int64_t device_buffer_samples = 0;

cpp/src/cluster/detail/kmeans.cuh

@@ -24,6 +24,7 @@
 #include <raft/core/pinned_mdarray.hpp>
 #include <raft/core/pinned_mdspan.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/linalg/map.cuh>
@@ -39,6 +40,7 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
+#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 
@@ -686,25 +688,62 @@ void kmeans_fit(
 
   auto minClusterAndDistance = raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(
     handle, device_buffer_samples);
-  auto L2NormBatch       = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
-  auto batch_weights_buf = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
+  auto minClusterDistance = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
+  auto L2NormBatch        = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
+  auto batch_weights_buf  = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
   rmm::device_uvector<DataT> L2NormBuf_OR_DistBuf(0, stream);
 
   auto centroid_sums      = raft::make_device_matrix<DataT, IndexT>(handle, n_clusters, n_features);
   auto weight_per_cluster = raft::make_device_vector<DataT, IndexT>(handle, n_clusters);
   auto clustering_cost    = raft::make_device_scalar<DataT>(handle, DataT{0});
+  auto batch_inertia      = raft::make_device_scalar<DataT>(handle, DataT{0});
 
   rmm::device_uvector<char> batch_workspace(device_buffer_samples, stream);
 
-  auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
-    handle, X.data_handle(), n_samples, n_features, device_buffer_samples, stream);
+  auto batch_mr = raft::resource::get_workspace_resource_ref(handle);
+
+  // Use the caller's stream pool for host-input staging. If no pool is configured, this falls
+  // back to the main stream and disables cross-stream prefetching.
+  auto [batch_copy_stream, enable_batch_prefetch] =
+    cuvs::spatial::knn::detail::utils::get_prefetch_stream(handle);
+
+  if constexpr (!data_on_device) {
+    size_t batch_staging_bytes =
+      static_cast<size_t>(device_buffer_samples) * static_cast<size_t>(n_features) * sizeof(DataT);
+    if (weight_ptr != nullptr) {
+      batch_staging_bytes += static_cast<size_t>(device_buffer_samples) * sizeof(DataT);
+    }
+    // Prefetch uses two staging buffers, so the workspace check must cover 2x the
+    // host-batch footprint (data and optional weights).
+    if (enable_batch_prefetch) { batch_staging_bytes *= 2; }
+    if (batch_staging_bytes > raft::resource::get_workspace_free_bytes(handle)) {
+      batch_mr = raft::resource::get_large_workspace_resource_ref(handle);
+    }
+  }
+
+  auto data_batches =
+    cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(handle,
+                                                                       X.data_handle(),
+                                                                       n_samples,
+                                                                       n_features,
+                                                                       device_buffer_samples,
+                                                                       batch_copy_stream,
+                                                                       batch_mr,
+                                                                       enable_batch_prefetch);
   // Host-path weight batches: only materialized when weights are provided and
   // the data resides on host
   std::optional<cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn<DataT>> weight_batches;
   if constexpr (!data_on_device) {
     if (weight_ptr != nullptr) {
-      weight_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
-        handle, weight_ptr, n_samples, IndexT{1}, device_buffer_samples, stream);
+      weight_batches =
+        cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(handle,
+                                                                           weight_ptr,
+                                                                           n_samples,
+                                                                           IndexT{1},
+                                                                           device_buffer_samples,
+                                                                           batch_copy_stream,
+                                                                           batch_mr,
+                                                                           enable_batch_prefetch);
     } else {
       raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1});
     }
@@ -836,19 +875,18 @@ void kmeans_fit(
         raft::make_device_matrix_view<DataT, IndexT>(new_centroids_ptr, n_clusters, n_features);
 
       data_batches.reset();
+      data_batches.prefetch_next_batch();
       using wt_iter_t = cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn<DataT>;
       std::optional<wt_iter_t> wt_it;
       if (weight_batches.has_value()) {
         weight_batches->reset();
         wt_it = weight_batches->begin();
+        wt_it->prefetch_next_batch();
       }
       for (const auto& data_batch : data_batches) {
         IndexT cur_batch_size = static_cast<IndexT>(data_batch.size());
         const DataT* wt_data  = nullptr;
-        if (wt_it.has_value()) {
-          wt_data = (**wt_it).data();
-          ++(*wt_it);
-        }
+        if (wt_it.has_value()) { wt_data = (**wt_it).data(); }
 
         auto batch_data_view = raft::make_device_matrix_view<const DataT, IndexT>(
           data_batch.data(), cur_batch_size, n_features);
@@ -893,6 +931,11 @@ void kmeans_fit(
                                      weight_per_cluster.view(),
                                      clustering_cost.view(),
                                      batch_workspace);
+        data_batches.prefetch_next_batch();
+        if (wt_it.has_value()) {
+          wt_it->prefetch_next_batch();
+          ++(*wt_it);
+        }
       }
       if (need_compute_norms) { norms_cached = true; }
 
@@ -933,40 +976,73 @@ void kmeans_fit(
       auto centroids_const = raft::make_device_matrix_view<const DataT, IndexT>(
         cur_centroids_ptr, n_clusters, n_features);
 
-      iter_inertia = DataT{0};
+      DataT zero = DataT{0};
+      raft::copy(clustering_cost.data_handle(), &zero, 1, stream);
       data_batches.reset();
+      data_batches.prefetch_next_batch();
       using wt_iter_t = cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn<DataT>;
       std::optional<wt_iter_t> wt_it;
       if (weight_batches.has_value()) {
         weight_batches->reset();
         wt_it = weight_batches->begin();
+        wt_it->prefetch_next_batch();
       }
       for (const auto& data_batch : data_batches) {
         IndexT cur_batch_size = static_cast<IndexT>(data_batch.size());
         const DataT* wt_data  = nullptr;
-        if (wt_it.has_value()) {
-          wt_data = (**wt_it).data();
-          ++(*wt_it);
-        }
+        if (wt_it.has_value()) { wt_data = (**wt_it).data(); }
 
         auto batch_data_view = raft::make_device_matrix_view<const DataT, IndexT>(
           data_batch.data(), cur_batch_size, n_features);
 
-        std::optional<raft::device_vector_view<const DataT, IndexT>> batch_sw = std::nullopt;
+        if constexpr (!data_on_device) {
+          if (need_compute_norms && norms_cached) {
+            raft::copy(L2NormBatch.data_handle(),
+                       h_norm_cache.data_handle() + data_batch.offset(),
+                       cur_batch_size,
+                       stream);
+          } else {
+            compute_batch_norms(data_batch.data(), cur_batch_size);
+          }
+        } else {
+          compute_batch_norms(data_batch.data(), cur_batch_size);
+        }
+        auto l2_norm_view =
+          raft::make_device_vector_view<DataT, IndexT>(L2NormBatch.data_handle(), cur_batch_size);
+        auto min_distance_view = raft::make_device_vector_view<DataT, IndexT>(
+          minClusterDistance.data_handle(), cur_batch_size);
+        std::optional<raft::device_vector_view<const DataT, IndexT>> batch_sample_weight =
+          std::nullopt;
         if (weight_ptr != nullptr) {
-          batch_sw =
+          batch_sample_weight =
             cur_batch_weights(static_cast<IndexT>(data_batch.offset()), wt_data, cur_batch_size);
         }
 
-        DataT batch_cost = DataT{0};
-        cuvs::cluster::kmeans::cluster_cost(handle,
-                                            batch_data_view,
-                                            centroids_const,
-                                            raft::make_host_scalar_view(&batch_cost),
-                                            batch_sw);
-
-        iter_inertia += batch_cost;
+        cluster_cost(handle,
+                     batch_data_view,
+                     centroids_const,
+                     min_distance_view,
+                     l2_norm_view,
+                     L2NormBuf_OR_DistBuf,
+                     cuvs::distance::DistanceType::L2Expanded,
+                     cur_batch_size,
+                     n_clusters,
+                     ws,
+                     batch_inertia.view(),
+                     batch_sample_weight);
+        raft::linalg::add(clustering_cost.data_handle(),
+                          clustering_cost.data_handle(),
+                          batch_inertia.data_handle(),
+                          1,
+                          stream);
+        data_batches.prefetch_next_batch();
+        if (wt_it.has_value()) {
+          wt_it->prefetch_next_batch();
+          ++(*wt_it);
+        }
       }
+      raft::copy(&iter_inertia, clustering_cost.data_handle(), 1, stream);
+      raft::resource::sync_stream(handle);
     }
 
     if (iter_inertia < inertia[0]) {

cpp/src/cluster/detail/kmeans_common.cuh

@@ -446,6 +446,50 @@ EXTERN_TEMPLATE_MIN_CLUSTER_DISTANCE(double, int)
 
 #undef EXTERN_TEMPLATE_MIN_CLUSTER_DISTANCE
 
+/**
+ * @brief Compute the optionally weighted sum of distances to the nearest centroid.
+ *
+ * Unlike minClusterAndDistanceCompute, this path does not calculate cluster labels.
+ */
+template <typename DataT, typename IndexT>
+void cluster_cost(
+  raft::resources const& handle,
+  raft::device_matrix_view<const DataT, IndexT> X,
+  raft::device_matrix_view<const DataT, IndexT> centroids,
+  raft::device_vector_view<DataT, IndexT> min_cluster_distance,
+  raft::device_vector_view<DataT, IndexT> l2_norm_x,
+  rmm::device_uvector<DataT>& l2_norm_or_distance_buffer,
+  cuvs::distance::DistanceType metric,
+  int batch_samples,
+  int batch_centroids,
+  rmm::device_uvector<char>& workspace,
+  raft::device_scalar_view<DataT> cost,
+  std::optional<raft::device_vector_view<const DataT, IndexT>> sample_weight = std::nullopt)
+{
+  auto centroids_mutable = raft::make_device_matrix_view<DataT, IndexT>(
+    const_cast<DataT*>(centroids.data_handle()), centroids.extent(0), centroids.extent(1));
+  minClusterDistanceCompute(handle,
+                            X,
+                            centroids_mutable,
+                            min_cluster_distance,
+                            l2_norm_x,
+                            l2_norm_or_distance_buffer,
+                            metric,
+                            batch_samples,
+                            batch_centroids,
+                            workspace);
+
+  if (sample_weight.has_value()) {
+    raft::linalg::map(handle,
+                      min_cluster_distance,
+                      raft::mul_op{},
+                      raft::make_const_mdspan(min_cluster_distance),
+                      sample_weight.value());
+  }
+  computeClusterCost(
+    handle, min_cluster_distance, workspace, cost, raft::identity_op{}, raft::add_op{});
+}
+
 template <typename DataT, typename IndexT>
 void countSamplesInCluster(raft::resources const& handle,
                            const cuvs::cluster::kmeans::params& params,

cpp/src/cluster/kmeans.cuh

@@ -343,7 +343,6 @@ void cluster_cost(
   auto stream     = raft::resource::get_cuda_stream(handle);
   auto n_clusters = centroids.extent(0);
   auto n_samples  = X.extent(0);
-  auto n_features = X.extent(1);
 
   rmm::device_uvector<char> workspace(n_samples * sizeof(IndexT), stream);
 
@@ -353,31 +352,18 @@ void cluster_cost(
   auto min_cluster_distance = raft::make_device_vector<DataT>(handle, n_samples);
   rmm::device_uvector<DataT> l2_norm_or_distance_buffer(0, stream);
 
-  auto metric = cuvs::distance::DistanceType::L2Expanded;
-
-  cuvs::cluster::kmeans::min_cluster_distance<DataT, IndexT>(
-    handle,
-    X,
-    raft::make_device_matrix_view<DataT, IndexT>(
-      const_cast<DataT*>(centroids.data_handle()), n_clusters, n_features),
-    min_cluster_distance.view(),
-    x_norms.view(),
-    l2_norm_or_distance_buffer,
-    metric,
-    n_samples,
-    n_clusters,
-    workspace);
-
-  if (sample_weight.has_value()) {
-    raft::linalg::map(handle,
-                      min_cluster_distance.view(),
-                      raft::mul_op{},
-                      raft::make_const_mdspan(min_cluster_distance.view()),
-                      sample_weight.value());
-  }
-
-  cuvs::cluster::kmeans::cluster_cost(
-    handle, min_cluster_distance.view(), workspace, cost, raft::add_op{});
+  cuvs::cluster::kmeans::detail::cluster_cost(handle,
+                                              X,
+                                              centroids,
+                                              min_cluster_distance.view(),
+                                              x_norms.view(),
+                                              l2_norm_or_distance_buffer,
+                                              cuvs::distance::DistanceType::L2Expanded,
+                                              n_samples,
+                                              n_clusters,
+                                              workspace,
+                                              cost,
+                                              sample_weight);
 }
 
 /**

python/cuvs/cuvs/cluster/kmeans/kmeans.pyx

@@ -89,6 +89,12 @@ cdef class KMeansParams:
         overhead as the number of times centroid adjustments are computed
         increases.
 
+        Host rows are staged through a device buffer of size
+        ``device_buffer_samples * n_features``. If the resources handle has a
+        CUDA stream pool with at least one stream, prefetch is enabled; that
+        doubles the batch memory footprint (and the weight-staging footprint
+        when sample weights are provided).
+
         Default: 0 (process all data at once).
     hierarchical : bool
         Whether to use hierarchical (balanced) kmeans or not