Exposing opt-in double buffering for MG batched KMeans

c/include/cuvs/cluster/kmeans.h

@@ -1,5 +1,5 @@
 /*
- * 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
  */
 
@@ -188,17 +188,25 @@ struct cuvsKMeansParams {
   int hierarchical_n_iters;
 
   /**
-   * 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).
+   * Number of host-resident samples to stage in a device buffer. When set to 0,
+   * defaults to n_samples (process all at once).
    */
-  int64_t streaming_batch_size;
+  int64_t device_buffer_batch_size;
 
   /**
    * Number of samples to draw for KMeansPlusPlus initialization.
    * When set to 0, uses heuristic min(3 * n_clusters, n_samples) for host data,
    * or n_samples for device data.
    */
   int64_t init_size;
+
+  /**
+   * Prefetches the next device buffer asynchronously while the current buffer
+   * is processed, hiding data-transfer latency for host-resident, multi-GPU
+   * KMeans. This overlaps transfers from host to device with computation using a
+   * second device batch buffer on each rank. Ignored by other KMeans paths.
+   */
+  bool device_buffer_prefetch;
 };
 
 typedef struct cuvsKMeansParams* cuvsKMeansParams_t;

c/src/cluster/kmeans.cpp

@@ -1,5 +1,5 @@
 /*
- * 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
  */
 
@@ -16,9 +16,8 @@
 
 namespace {
 
-// The conversions are templated on the C struct type and reused by both API surfaces.
-template <typename ParamsT>
-cuvs::cluster::kmeans::params convert_params(const ParamsT& params)
+// The legacy C struct keeps its original field names for ABI compatibility.
+cuvs::cluster::kmeans::params convert_params(const cuvsKMeansParams& params)
 {
   auto kmeans_params                = cuvs::cluster::kmeans::params();
   kmeans_params.metric              = static_cast<cuvs::distance::DistanceType>(params.metric);
@@ -31,7 +30,24 @@ cuvs::cluster::kmeans::params convert_params(const ParamsT& params)
   kmeans_params.batch_samples       = params.batch_samples;
   kmeans_params.batch_centroids     = params.batch_centroids;
   kmeans_params.init_size             = params.init_size;
-  kmeans_params.streaming_batch_size  = params.streaming_batch_size;
+  kmeans_params.device_buffer_batch_size = params.streaming_batch_size;
+  return kmeans_params;
+}
+
+cuvs::cluster::kmeans::params convert_params(const cuvsKMeansParams_v2& params)
+{
+  auto kmeans_params                = cuvs::cluster::kmeans::params();
+  kmeans_params.metric              = static_cast<cuvs::distance::DistanceType>(params.metric);
+  kmeans_params.init = static_cast<cuvs::cluster::kmeans::params::InitMethod>(params.init);
+  kmeans_params.n_clusters          = params.n_clusters;
+  kmeans_params.max_iter            = params.max_iter;
+  kmeans_params.tol                 = params.tol;
+  kmeans_params.n_init              = params.n_init;
+  kmeans_params.oversampling_factor = params.oversampling_factor;
+  kmeans_params.batch_samples       = params.batch_samples;
+  kmeans_params.batch_centroids     = params.batch_centroids;
+  kmeans_params.init_size             = params.init_size;
+  kmeans_params.device_buffer_batch_size = params.device_buffer_batch_size;
   return kmeans_params;
 }
 
@@ -243,7 +259,7 @@ extern "C" cuvsError_t cuvsKMeansParamsCreate(cuvsKMeansParams_t* params)
       .inertia_check        = false,
       .hierarchical         = false,
       .hierarchical_n_iters = static_cast<int>(cpp_balanced_params.n_iters),
-      .streaming_batch_size = cpp_params.streaming_batch_size,
+      .streaming_batch_size = cpp_params.device_buffer_batch_size,
       .init_size            = cpp_params.init_size};
   });
 }
@@ -315,8 +331,9 @@ extern "C" cuvsError_t cuvsKMeansParamsCreate_v2(cuvsKMeansParams_v2_t* params)
       .batch_centroids      = cpp_params.batch_centroids,
       .hierarchical         = false,
       .hierarchical_n_iters = static_cast<int>(cpp_balanced_params.n_iters),
-      .streaming_batch_size = cpp_params.streaming_batch_size,
-      .init_size            = cpp_params.init_size};
+      .device_buffer_batch_size = cpp_params.device_buffer_batch_size,
+      .init_size            = cpp_params.init_size,
+      .device_buffer_prefetch = cpp_params.device_buffer_prefetch};
   });
 }
 

c/src/cluster/mg_kmeans.cpp

@@ -37,7 +37,8 @@ cuvs::cluster::kmeans::params convert_params(const ParamsT& params)
   kmeans_params.batch_samples        = params.batch_samples;
   kmeans_params.batch_centroids      = params.batch_centroids;
   kmeans_params.init_size            = params.init_size;
-  kmeans_params.streaming_batch_size = params.streaming_batch_size;
+  kmeans_params.device_buffer_batch_size = params.device_buffer_batch_size;
+  kmeans_params.device_buffer_prefetch = params.device_buffer_prefetch;
   return kmeans_params;
 }
 

c/tests/cluster/kmeans_mg_c.cu

@@ -73,11 +73,13 @@ void test_mg_fit_host()
 
   typename Api::params_t params;
   ASSERT_EQ(Api::params_create(&params), CUVS_SUCCESS);
+  EXPECT_FALSE(params->device_buffer_prefetch);
   params->n_clusters           = kNClusters;
   params->max_iter             = 100;
   params->tol                  = 1e-6;
   params->init                 = Array;
-  params->streaming_batch_size = 4;  // force at least 2 streamed batches
+  params->device_buffer_batch_size = 4;  // force at least two device buffers
+  params->device_buffer_prefetch = true;
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(raft::make_host_matrix_view<float, int64_t>(

cpp/include/cuvs/cluster/kmeans.hpp

@@ -113,7 +113,7 @@ struct params : base_params {
    * Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0
    * then don't tile the centroids
    *
-   * NB: These parameters are unrelated to streaming_batch_size, which controls how many
+   * NB: These parameters are unrelated to device_buffer_batch_size, which controls how many
    * samples to transfer from host to device per batch when processing out-of-core
    * data.
    */
@@ -144,17 +144,28 @@ struct params : base_params {
   int64_t init_size = 0;
 
   /**
-   * Number of samples to process per GPU batch when fitting with host data.
+   * Number of host-resident samples staged in each device buffer.
    * When set to 0, defaults to n_samples (process all at once).
    * Only used by the batched (host-data) code path and ignored by
    * device-data overloads.
    *
-   * 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
+   * In multi-GPU mode this is a per-rank buffer size: each rank processes up
+   * to this many local samples at once, clamped to that rank's local sample
    * count. This is is ignored by device-data overloads.
    * Default: 0 (process all data at once).
    */
-  int64_t streaming_batch_size = 0;
+  int64_t device_buffer_batch_size = 0;
+
+  /**
+   * Prefetches the next device buffer asynchronously on a separate CUDA stream
+   * while the current buffer is processed, hiding data-transfer latency for
+   * host-resident, multi-GPU KMeans. This overlaps transfers from host to device
+   * with computation using a second device buffer on each rank.
+   *
+   * This option is ignored by single-GPU and device-resident fits.
+   * Default: false.
+   */
+  bool device_buffer_prefetch = false;
 };
 
 /**
@@ -199,7 +210,7 @@ enum class kmeans_type { KMeans = 0, KMeansBalanced = 1 };
  *
  * This overload supports out-of-core computation where the dataset resides
  * on the host. Data is processed in batches, streaming from host to
- * device. The batch size is controlled by `params.streaming_batch_size`.
+ * device. The batch size is controlled by `params.device_buffer_batch_size`.
  *
  * Multi-GPU dispatch is selected automatically based on the handle state:
  *   - If `raft::resource::is_multi_gpu(handle)` (cuVS SNMG): the full dataset X
@@ -221,7 +232,7 @@ enum class kmeans_type { KMeans = 0, KMeansBalanced = 1 };
  *   raft::resources handle;
  *   cuvs::cluster::kmeans::params params;
  *   params.n_clusters = 100;
- *   params.streaming_batch_size = 100000;
+ *   params.device_buffer_batch_size = 100000;
  *   float inertia;
  *   int64_t n_iter;
  *
@@ -245,7 +256,7 @@ enum class kmeans_type { KMeans = 0, KMeansBalanced = 1 };
  * @param[in]     handle        The raft handle. When a multi-GPU resource is
  *                              attached, multi-GPU dispatch is used automatically.
  * @param[in]     params        Parameters for KMeans model. Batch size is read from
- *                              params.streaming_batch_size.
+ *                              params.device_buffer_batch_size.
  * @param[in]     X             Training instances on HOST memory. The data must
  *                              be in row-major format.
  *                              [dim = n_samples x n_features]
@@ -1648,8 +1659,8 @@ void cluster_cost(
  *
  * Each rank supplies its local training data as a vector of partitions. For
  * host-resident partitions the implementation streams each partition using
- * `params.streaming_batch_size` (per rank). For device-resident partitions
- * `streaming_batch_size` is ignored and each local partition is processed in full.
+ * `params.device_buffer_batch_size` (per rank). For device-resident partitions
+ * `device_buffer_batch_size` is ignored and each local partition is processed in full.
  *
  * The active backend is selected by the resources attached to
  * `handle`:
@@ -1663,9 +1674,9 @@ void cluster_cost(
  * @param[in]     handle              The raft handle. Must have NCCL comms or
  *                                    a SNMG clique initialized.
  * @param[in]     params              K-means parameters. For host-resident
- *                                    partitions the per-rank streaming batch
+ *                                    partitions the per-rank device-buffer
  *                                    size is read from
- *                                    `params.streaming_batch_size`; it is
+ *                                    `params.device_buffer_batch_size`; it is
  *                                    ignored for device-resident partitions.
  * @param[in]     X_parts             Per-partition local data on this rank.
  *                                    Each entry is [n_rows_i x n_features].

cpp/src/cluster/detail/kmeans.cuh

@@ -588,9 +588,9 @@ void kmeans_fit(
 
   raft::default_logger().set_level(pams.verbosity);
 
-  IndexT streaming_batch_size = static_cast<IndexT>(pams.streaming_batch_size);
-  if (streaming_batch_size <= 0 || streaming_batch_size > static_cast<IndexT>(n_samples)) {
-    streaming_batch_size = static_cast<IndexT>(n_samples);
+  IndexT device_buffer_batch_size = static_cast<IndexT>(pams.device_buffer_batch_size);
+  if (device_buffer_batch_size <= 0 || device_buffer_batch_size > static_cast<IndexT>(n_samples)) {
+    device_buffer_batch_size = static_cast<IndexT>(n_samples);
   }
 
   constexpr bool data_on_device = raft::is_device_mdspan_v<decltype(X)>;
@@ -606,13 +606,13 @@ void kmeans_fit(
   rmm::device_uvector<char> local_workspace(0, stream);
   rmm::device_uvector<char>& ws = workspace.has_value() ? workspace->get() : local_workspace;
 
-  if (data_on_device && streaming_batch_size != static_cast<IndexT>(n_samples)) {
+  if (data_on_device && device_buffer_batch_size != static_cast<IndexT>(n_samples)) {
     RAFT_LOG_WARN(
-      "KMeans: streaming_batch_size (%zu) ignored when data resides on device; using n_samples "
+      "KMeans: device_buffer_batch_size (%zu) ignored when data resides on device; using n_samples "
       "(%zu)",
-      static_cast<size_t>(streaming_batch_size),
+      static_cast<size_t>(device_buffer_batch_size),
       static_cast<size_t>(n_samples));
-    streaming_batch_size = static_cast<IndexT>(n_samples);
+    device_buffer_batch_size = static_cast<IndexT>(n_samples);
   }
 
   // Preallocate the host-side KMeans++ init sample buffer.
@@ -685,26 +685,27 @@ void kmeans_fit(
   DataT* new_centroids_ptr = new_centroids_buf.data();
 
   auto minClusterAndDistance = raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(
-    handle, streaming_batch_size);
-  auto L2NormBatch       = raft::make_device_vector<DataT, IndexT>(handle, streaming_batch_size);
-  auto batch_weights_buf = raft::make_device_vector<DataT, IndexT>(handle, streaming_batch_size);
+    handle, device_buffer_batch_size);
+  auto L2NormBatch = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_batch_size);
+  auto batch_weights_buf =
+    raft::make_device_vector<DataT, IndexT>(handle, device_buffer_batch_size);
   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});
 
-  rmm::device_uvector<char> batch_workspace(streaming_batch_size, stream);
+  rmm::device_uvector<char> batch_workspace(device_buffer_batch_size, stream);
 
   auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
-    handle, X.data_handle(), n_samples, n_features, streaming_batch_size, stream);
+    handle, X.data_handle(), n_samples, n_features, device_buffer_batch_size, stream);
   // 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}, streaming_batch_size, stream);
+        handle, weight_ptr, n_samples, IndexT{1}, device_buffer_batch_size, stream);
     } else {
       raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1});
     }
@@ -759,11 +760,11 @@ void kmeans_fit(
   };
 
   RAFT_LOG_DEBUG(
-    "KMeans.fit: n_samples=%zu, n_features=%zu, n_clusters=%d, streaming_batch_size=%zu",
+    "KMeans.fit: n_samples=%zu, n_features=%zu, n_clusters=%d, device_buffer_batch_size=%zu",
     static_cast<size_t>(n_samples),
     static_cast<size_t>(n_features),
     n_clusters,
-    static_cast<size_t>(streaming_batch_size));
+    static_cast<size_t>(device_buffer_batch_size));
 
   bool need_compute_norms = metric == cuvs::distance::DistanceType::L2Expanded ||
                             metric == cuvs::distance::DistanceType::L2SqrtExpanded;