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
  */
 
@@ -199,6 +199,12 @@ struct cuvsKMeansParams {
    * or n_samples for device data.
    */
   int64_t init_size;
+
+  /**
+   * Whether host-resident multi-GPU KMeans should prefetch the next streaming
+   * batch using a second device buffer. Ignored by other KMeans paths.
+   */
+  bool streaming_batch_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
  */
 
@@ -316,7 +316,8 @@ extern "C" cuvsError_t cuvsKMeansParamsCreate_v2(cuvsKMeansParams_v2_t* params)
       .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};
+      .init_size            = cpp_params.init_size,
+      .streaming_batch_prefetch = cpp_params.streaming_batch_prefetch};
   });
 }
 

c/src/cluster/mg_kmeans.cpp

@@ -38,6 +38,7 @@ cuvs::cluster::kmeans::params convert_params(const ParamsT& params)
   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.streaming_batch_prefetch = params.streaming_batch_prefetch;
   return kmeans_params;
 }
 

c/tests/cluster/kmeans_c.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -54,6 +54,10 @@ struct kmeans_api_v1 {
   using params_t = cuvsKMeansParams_t;
   static cuvsError_t params_create(params_t* p) { return cuvsKMeansParamsCreate(p); }
   static cuvsError_t params_destroy(params_t p) { return cuvsKMeansParamsDestroy(p); }
+  static void set_device_buffer_batch_size(params_t params, int64_t batch_size)
+  {
+    params->streaming_batch_size = batch_size;
+  }
   static cuvsError_t fit(cuvsResources_t res,
                          params_t params,
                          DLManagedTensor* dataset,
@@ -79,6 +83,10 @@ struct kmeans_api_v2 {
   using params_t = cuvsKMeansParams_v2_t;
   static cuvsError_t params_create(params_t* p) { return cuvsKMeansParamsCreate_v2(p); }
   static cuvsError_t params_destroy(params_t p) { return cuvsKMeansParamsDestroy_v2(p); }
+  static void set_device_buffer_batch_size(params_t params, int64_t batch_size)
+  {
+    params->device_buffer_batch_size = batch_size;
+  }
   static cuvsError_t fit(cuvsResources_t res,
                          params_t params,
                          DLManagedTensor* dataset,
@@ -128,7 +136,7 @@ void test_fit_predict()
   params->max_iter             = 100;
   params->tol                  = 1e-6;
   params->init                 = Array;
-  params->streaming_batch_size = 0;
+  Api::set_device_buffer_batch_size(params, 0);
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(
@@ -195,7 +203,7 @@ void test_fit_host()
   params->max_iter             = 100;
   params->tol                  = 1e-6;
   params->init                 = Array;
-  params->streaming_batch_size = 4;  // force at least 2 streamed batches
+  Api::set_device_buffer_batch_size(params, 4);  // force at least two device buffers
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(

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->streaming_batch_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->streaming_batch_prefetch = true;
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(raft::make_host_matrix_view<float, int64_t>(

cpp/include/cuvs/cluster/kmeans.hpp

@@ -155,6 +155,16 @@ struct params : base_params {
    * Default: 0 (process all data at once).
    */
   int64_t streaming_batch_size = 0;
+
+  /**
+   * Whether host-resident multi-GPU KMeans should prefetch the next streaming
+   * batch on a separate CUDA stream. Enabling this overlaps H2D transfer with
+   * computation by allocating a second device batch buffer on each rank.
+   *
+   * This option is ignored by single-GPU and device-resident fits.
+   * Default: false.
+   */
+  bool streaming_batch_prefetch = false;
 };
 
 /**

cpp/src/cluster/detail/kmeans_mg.cuh

@@ -37,13 +37,15 @@
 #include <raft/random/rng.cuh>
 #include <raft/util/cudart_utils.hpp>
 
+#include <rmm/cuda_stream.hpp>
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <algorithm>
 #include <cmath>
 #include <cstdint>
 #include <limits>
+#include <memory>
 #include <numeric>
 #include <optional>
 #include <random>
@@ -140,7 +142,18 @@ void mnmg_fit(
     use_nccl ? raft::resource::set_current_device_to_rank(handle, rank) : handle;
   mnmg_comms comms{dev_res, use_nccl, nccl_comm};
 
-  auto stream     = comms.stream();
+  auto stream = comms.stream();
+  std::unique_ptr<rmm::cuda_stream> data_copy_stream_owner;
+  rmm::cuda_stream_view data_copy_stream{stream};
+  bool enable_data_prefetch = false;
+  if constexpr (!data_on_device) {
+    if (params.streaming_batch_prefetch) {
+      data_copy_stream_owner =
+        std::make_unique<rmm::cuda_stream>(rmm::cuda_stream::flags::non_blocking);
+      data_copy_stream     = data_copy_stream_owner->view();
+      enable_data_prefetch = true;
+    }
+  }
   auto n_features = centroids.extent(1);
   auto n_clusters = static_cast<IndexT>(params.n_clusters);
   auto metric     = params.metric;
@@ -354,7 +367,25 @@ void mnmg_fit(
   };
   auto rank_centroids_const = raft::make_const_mdspan(rank_centroids);
 
+  // Persist one-partition input across Lloyd iterations and inertia. Multi-partition inputs keep
+  // the bounded-memory transient path.
+  const bool persist_data_batches = X_parts.size() == 1 && X_parts.front().extent(0) > 0;
+  std::optional<data_batch_iterator_t> persistent_data_batches;
+  auto make_data_batches = [&](const data_part_view_t& X_part) -> data_batch_iterator_t {
+    return persistent_data_batches
+             ? *persistent_data_batches
+             : data_batch_iterator_t(dev_res,
+                                     X_part,
+                                     static_cast<size_t>(streaming_batch_size),
+                                     data_copy_stream,
+                                     rmm::mr::get_current_device_resource_ref(),
+                                     enable_data_prefetch);
+  };
+
   for (int seed_iter = 0; seed_iter < n_init; ++seed_iter) {
+    // Free the persistent batch during initialization, then retain it for Lloyd and inertia.
+    persistent_data_batches.reset();
+
     cuvs::cluster::kmeans::params iter_params = params;
     iter_params.rng_state.seed                = gen();
 
@@ -374,6 +405,15 @@ void mnmg_fit(
                                                            rank,
                                                            comms);
 
+    if (persist_data_batches) {
+      persistent_data_batches.emplace(dev_res,
+                                      X_parts.front(),
+                                      static_cast<size_t>(streaming_batch_size),
+                                      data_copy_stream,
+                                      rmm::mr::get_current_device_resource_ref(),
+                                      enable_data_prefetch);
+    }
+
     if (!sample_weights) { raft::matrix::fill(dev_res, batch_weights.view(), DataT{1}); }
 
     raft::matrix::fill(dev_res, d_prior_cost.view(), DataT{0});
@@ -402,14 +442,13 @@ void mnmg_fit(
         auto part_rows     = static_cast<IndexT>(X_part.extent(0));
         if (part_rows == 0) { continue; }
 
-        data_batch_iterator_t data_batches(dev_res,
-                                           X_part,
-                                           static_cast<size_t>(streaming_batch_size),
-                                           stream,
-                                           rmm::mr::get_current_device_resource_ref(),
-                                           true);
+        auto data_batches = make_data_batches(X_part);
+        auto data_it      = data_batches.begin();
+        auto data_end     = data_batches.end();
+        data_batches.prefetch_next_batch();
 
-        for (auto const& data_batch : data_batches) {
+        for (; data_it != data_end; ++data_it) {
+          auto const& data_batch    = *data_it;
           IndexT current_batch_size = static_cast<IndexT>(data_batch.size());
           auto batch_offset         = static_cast<IndexT>(data_batch.offset());
 
@@ -468,7 +507,9 @@ void mnmg_fit(
             weight_per_cluster.view(),
             raft::make_device_scalar_view(clustering_cost.data_handle()),
             batch_workspace);
+          data_batches.prefetch_next_batch();
         }
+        if (enable_data_prefetch) { raft::resource::sync_stream(dev_res); }
       }
       norms_cached = true;
 
@@ -530,14 +571,13 @@ void mnmg_fit(
       auto part_rows     = static_cast<IndexT>(X_part.extent(0));
       if (part_rows == 0) { continue; }
 
-      data_batch_iterator_t data_batches(dev_res,
-                                         X_part,
-                                         static_cast<size_t>(streaming_batch_size),
-                                         stream,
-                                         rmm::mr::get_current_device_resource_ref(),
-                                         true);
+      auto data_batches = make_data_batches(X_part);
+      auto data_it      = data_batches.begin();
+      auto data_end     = data_batches.end();
+      data_batches.prefetch_next_batch();
 
-      for (auto const& data_batch : data_batches) {
+      for (; data_it != data_end; ++data_it) {
+        auto const& data_batch    = *data_it;
         IndexT current_batch_size = static_cast<IndexT>(data_batch.size());
         auto batch_offset         = static_cast<IndexT>(data_batch.offset());
 
@@ -561,7 +601,9 @@ void mnmg_fit(
                           raft::make_const_mdspan(clustering_cost.view()),
                           raft::make_const_mdspan(batch_clustering_cost.view()),
                           clustering_cost.view());
+        data_batches.prefetch_next_batch();
       }
+      if (enable_data_prefetch) { raft::resource::sync_stream(dev_res); }
     }
     comms.allreduce(clustering_cost.data_handle(), clustering_cost.data_handle(), 1);
     raft::copy(&local_inertia, clustering_cost.data_handle(), 1, stream);

cpp/tests/cluster/kmeans_mg.cu

@@ -42,6 +42,7 @@ struct KmeansSNMGInputs {
   int n_init;
   cuvs::cluster::kmeans::params::InitMethod init = cuvs::cluster::kmeans::params::Array;
   int max_iter                                   = 20;
+  bool streaming_batch_prefetch                  = false;
 };
 
 template <typename T>
@@ -105,13 +106,14 @@ class KmeansSNMGTest : public ::testing::TestWithParam<KmeansSNMGInputs<T>> {
     }
 
     cuvs::cluster::kmeans::params snmg_params;
-    snmg_params.n_clusters           = n_clusters;
-    snmg_params.tol                  = testparams_.tol;
-    snmg_params.max_iter             = testparams_.max_iter;
-    snmg_params.n_init               = testparams_.n_init;
-    snmg_params.rng_state.seed       = 42;
-    snmg_params.init                 = testparams_.init;
-    snmg_params.streaming_batch_size = testparams_.streaming_batch_size;
+    snmg_params.n_clusters               = n_clusters;
+    snmg_params.tol                      = testparams_.tol;
+    snmg_params.max_iter                 = testparams_.max_iter;
+    snmg_params.n_init                   = testparams_.n_init;
+    snmg_params.rng_state.seed           = 42;
+    snmg_params.init                     = testparams_.init;
+    snmg_params.streaming_batch_size     = testparams_.streaming_batch_size;
+    snmg_params.streaming_batch_prefetch = testparams_.streaming_batch_prefetch;
 
     T snmg_inertia      = T{0};
     int64_t snmg_n_iter = 0;
@@ -281,6 +283,28 @@ const std::vector<KmeansSNMGInputs<float>> snmg_inputsf = {
   {1000, 32, 5, 0.0001f, kmeans_weight_mode::none, 1000, 1},
   {1000, 32, 5, 0.0001f, kmeans_weight_mode::uniform, 1000, 1},
   {1000, 32, 5, 0.0001f, kmeans_weight_mode::none, 128, 1},
+  // Prefetch enabled with a final partial device buffer (1000 % 128 != 0).
+  {1000,
+   32,
+   5,
+   0.0001f,
+   kmeans_weight_mode::none,
+   128,
+   1,
+   cuvs::cluster::kmeans::params::Array,
+   20,
+   true},
+  // Prefetch enabled with one device buffer: the iterator must use its single-buffer fast path.
+  {1000,
+   32,
+   5,
+   0.0001f,
+   kmeans_weight_mode::none,
+   1000,
+   1,
+   cuvs::cluster::kmeans::params::Array,
+   20,
+   true},
   {10000, 16, 10, 0.0001f, kmeans_weight_mode::none, 2000, 1},
   {10000, 16, 10, 0.0001f, kmeans_weight_mode::uniform, 2000, 1},
   {10000, 16, 10, 0.0001f, kmeans_weight_mode::none, 500, 1},

python/cuvs/cuvs/cluster/kmeans/kmeans.pxd

@@ -56,7 +56,8 @@ cdef extern from "cuvs/cluster/kmeans.h" nogil:
         bool hierarchical,
         int hierarchical_n_iters,
         int64_t streaming_batch_size,
-        int64_t init_size
+        int64_t init_size,
+        bool streaming_batch_prefetch
 
     ctypedef cuvsKMeansParams* cuvsKMeansParams_t
     ctypedef cuvsKMeansParams_v2* cuvsKMeansParams_v2_t
@@ -90,3 +91,4 @@ cdef extern from "cuvs/cluster/kmeans.h" nogil:
 
 cdef class KMeansParams:
     cdef cuvsKMeansParams* params
+    cdef bool _streaming_batch_prefetch

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

@@ -94,13 +94,21 @@ cdef class KMeansParams:
         increases.
 
         Default: 0 (process all data at once).
+    streaming_batch_prefetch : bool
+        Whether host-resident multi-GPU KMeans should use a second device
+        buffer to prefetch the next streaming batch and overlap H2D transfer
+        with computation. This can improve throughput at the cost of one
+        additional device batch buffer per GPU. Ignored by other KMeans paths.
+
+        Default: False.
     hierarchical : bool
         Whether to use hierarchical (balanced) kmeans or not
     hierarchical_n_iters : int
         For hierarchical k-means , defines the number of training iterations
     """
 
     def __cinit__(self):
+        self._streaming_batch_prefetch = False
         cuvsKMeansParamsCreate(&self.params)
 
     def __dealloc__(self):
@@ -119,6 +127,7 @@ cdef class KMeansParams:
                  inertia_check=None,
                  init_size=None,
                  streaming_batch_size=None,
+                 streaming_batch_prefetch=None,
                  hierarchical=None,
                  hierarchical_n_iters=None):
         if metric is not None:
@@ -150,6 +159,8 @@ cdef class KMeansParams:
             self.params.init_size = init_size
         if streaming_batch_size is not None:
             self.params.streaming_batch_size = streaming_batch_size
+        if streaming_batch_prefetch is not None:
+            self._streaming_batch_prefetch = streaming_batch_prefetch
         if hierarchical is not None:
             self.params.hierarchical = hierarchical
         if hierarchical_n_iters is not None:
@@ -202,6 +213,10 @@ cdef class KMeansParams:
     def streaming_batch_size(self):
         return self.params.streaming_batch_size
 
+    @property
+    def streaming_batch_prefetch(self):
+        return self._streaming_batch_prefetch
+
     @property
     def hierarchical(self):
         return self.params.hierarchical