Merge pull request #2356 from NVIDIA/release/26.08

Forward-merge release/26.08 into main

Author: GPUtester

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
  */
 
@@ -112,7 +112,7 @@ 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).
    */
-  int64_t streaming_batch_size;
+  int64_t device_buffer_samples;
 
   /**
    * Number of samples to draw for KMeansPlusPlus initialization.
@@ -191,7 +191,7 @@ 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).
    */
-  int64_t streaming_batch_size;
+  int64_t device_buffer_samples;
 
   /**
    * Number of samples to draw for KMeansPlusPlus initialization.
@@ -267,8 +267,8 @@ typedef enum { CUVS_KMEANS_TYPE_KMEANS = 0, CUVS_KMEANS_TYPE_KMEANS_BALANCED = 1
  *   k-means++ algorithm.
  *
  *   X may reside on either host (CPU) or device (GPU) memory.
- *   When X is on the host the data is streamed to the GPU in
- *   batches controlled by params->streaming_batch_size.
+ *   When X is on the host the data is buffered to the GPU in
+ *   batches controlled by params->device_buffer_samples.
  *
  * @note In cuVS 26.08 (next ABI major version) this signature will be
  * replaced by cuvsKMeansFit_v2.

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
  */
 
@@ -31,7 +31,7 @@ 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_samples  = params.device_buffer_samples;
   return kmeans_params;
 }
 
@@ -243,7 +243,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,
+      .device_buffer_samples = cpp_params.device_buffer_samples,
       .init_size            = cpp_params.init_size};
   });
 }
@@ -315,7 +315,7 @@ 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,
+      .device_buffer_samples = cpp_params.device_buffer_samples,
       .init_size            = cpp_params.init_size};
   });
 }

c/src/cluster/mg_kmeans.cpp

@@ -37,7 +37,7 @@ 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_samples = params.device_buffer_samples;
   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
  */
 
@@ -128,7 +128,7 @@ void test_fit_predict()
   params->max_iter             = 100;
   params->tol                  = 1e-6;
   params->init                 = Array;
-  params->streaming_batch_size = 0;
+  params->device_buffer_samples = 0;
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(
@@ -195,7 +195,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
+  params->device_buffer_samples = 4;  // force at least 2 streamed batches
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(

c/tests/cluster/kmeans_mg_c.cu

@@ -77,7 +77,7 @@ void test_mg_fit_host()
   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_samples = 4;  // force at least 2 streamed batches
 
   DLManagedTensor dataset_t{};
   cuvs::core::to_dlpack(raft::make_host_matrix_view<float, int64_t>(

cpp/include/cuvs/cluster/kmeans.hpp

@@ -113,9 +113,8 @@ 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
-   * samples to transfer from host to device per batch when processing out-of-core
-   * data.
+   * NB: These parameters are unrelated to device_buffer_samples, which specifies the number of
+   * training vectors that get buffered on device when the training vectors are passed in on host.
    */
   int batch_samples = 1 << 15;
 
@@ -154,7 +153,7 @@ struct params : base_params {
    * 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_samples = 0;
 };
 
 /**
@@ -239,7 +238,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_samples`.
  *
  * Multi-GPU dispatch is selected automatically based on the handle state:
  *   - If `raft::resource::is_multi_gpu(handle)` (cuVS SNMG): the full dataset X
@@ -261,7 +260,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_samples = 100000;
  *   float inertia;
  *   int64_t n_iter;
  *
@@ -285,7 +284,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_samples.
  * @param[in]     X             Training instances on HOST memory. The data must
  *                              be in row-major format.
  *                              [dim = n_samples x n_features]
@@ -1688,8 +1687,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_samples` (per rank). For device-resident partitions
+ * `device_buffer_samples` is ignored and each local partition is processed in full.
  *
  * The active backend is selected by the resources attached to
  * `handle`:
@@ -1705,7 +1704,7 @@ void cluster_cost(
  * @param[in]     params              K-means parameters. For host-resident
  *                                    partitions the per-rank streaming batch
  *                                    size is read from
- *                                    `params.streaming_batch_size`; it is
+ *                                    `params.device_buffer_samples`; 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_samples = static_cast<IndexT>(pams.device_buffer_samples);
+  if (device_buffer_samples <= 0 || device_buffer_samples > static_cast<IndexT>(n_samples)) {
+    device_buffer_samples = 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_samples != 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_samples (%zu) ignored when data resides on device; using n_samples "
       "(%zu)",
-      static_cast<size_t>(streaming_batch_size),
+      static_cast<size_t>(device_buffer_samples),
       static_cast<size_t>(n_samples));
-    streaming_batch_size = static_cast<IndexT>(n_samples);
+    device_buffer_samples = static_cast<IndexT>(n_samples);
   }
 
   // Preallocate the host-side KMeans++ init sample buffer.
@@ -685,26 +685,26 @@ 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_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});
 
-  rmm::device_uvector<char> batch_workspace(streaming_batch_size, stream);
+  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, streaming_batch_size, stream);
+    handle, X.data_handle(), n_samples, n_features, device_buffer_samples, 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_samples, stream);
     } else {
       raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1});
     }
@@ -759,11 +759,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_samples=%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_samples));
 
   bool need_compute_norms = metric == cuvs::distance::DistanceType::L2Expanded ||
                             metric == cuvs::distance::DistanceType::L2SqrtExpanded;

cpp/src/cluster/detail/kmeans_mg.cuh

@@ -188,18 +188,18 @@ void mnmg_fit(
                  static_cast<size_t>(n_features),
                  static_cast<int>(n_clusters));
 
-  IndexT streaming_batch_size = static_cast<IndexT>(params.streaming_batch_size);
-  if (streaming_batch_size <= 0 || streaming_batch_size > max_part_rows) {
-    streaming_batch_size = std::max(max_part_rows, IndexT{1});
+  IndexT device_buffer_samples = static_cast<IndexT>(params.device_buffer_samples);
+  if (device_buffer_samples <= 0 || device_buffer_samples > max_part_rows) {
+    device_buffer_samples = std::max(max_part_rows, IndexT{1});
   }
 
-  if (data_on_device && streaming_batch_size < max_part_rows) {
+  if (data_on_device && device_buffer_samples < max_part_rows) {
     RAFT_LOG_WARN(
-      "MNMG KMeans: streaming_batch_size (%zu) ignored when partitions reside on device; using "
+      "MNMG KMeans: device_buffer_samples (%zu) ignored when partitions reside on device; using "
       "max partition size (%zu)",
-      static_cast<size_t>(streaming_batch_size),
+      static_cast<size_t>(device_buffer_samples),
       static_cast<size_t>(max_part_rows));
-    streaming_batch_size = max_part_rows;
+    device_buffer_samples = max_part_rows;
   }
 
   auto rank_centroids_arr =
@@ -212,7 +212,7 @@ void mnmg_fit(
   auto clustering_cost       = raft::make_device_vector<DataT, IndexT>(dev_res, 1);
   auto batch_clustering_cost = raft::make_device_vector<DataT, IndexT>(dev_res, 1);
   auto sqrd_norm_error_dev   = raft::make_device_scalar<DataT>(dev_res, DataT{0});
-  IndexT alloc_batch_size    = streaming_batch_size;
+  IndexT alloc_batch_size    = device_buffer_samples;
   auto batch_weights         = raft::make_device_vector<DataT, IndexT>(dev_res, alloc_batch_size);
   auto minClusterAndDistance =
     raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(dev_res, alloc_batch_size);
@@ -363,7 +363,7 @@ void mnmg_fit(
 
     init_centroids_for_mg_batched<DataT, IndexT, Accessor>(dev_res,
                                                            iter_params,
-                                                           streaming_batch_size,
+                                                           device_buffer_samples,
                                                            X_parts,
                                                            n_features,
                                                            input_centroids_const,
@@ -404,7 +404,7 @@ void mnmg_fit(
 
         data_batch_iterator_t data_batches(dev_res,
                                            X_part,
-                                           static_cast<size_t>(streaming_batch_size),
+                                           static_cast<size_t>(device_buffer_samples),
                                            stream,
                                            rmm::mr::get_current_device_resource_ref(),
                                            true);
@@ -532,7 +532,7 @@ void mnmg_fit(
 
       data_batch_iterator_t data_batches(dev_res,
                                          X_part,
-                                         static_cast<size_t>(streaming_batch_size),
+                                         static_cast<size_t>(device_buffer_samples),
                                          stream,
                                          rmm::mr::get_current_device_resource_ref(),
                                          true);

cpp/src/cluster/detail/kmeans_mg_batched_init.cuh

@@ -247,7 +247,7 @@ template <typename DataT, typename IndexT, typename Accessor>
 void init_centroids_for_mg_batched(
   raft::resources const& handle,
   const cuvs::cluster::kmeans::params& params,
-  IndexT /*streaming_batch_size*/,
+  IndexT /*device_buffer_samples*/,
   const std::vector<
     raft::mdspan<const DataT, raft::matrix_extent<IndexT>, raft::row_major, Accessor>>& X_parts,
   IndexT n_features,

cpp/tests/cluster/kmeans.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -339,7 +339,7 @@ struct KmeansBatchedInputs {
   int n_clusters;
   T tol;
   bool weighted;
-  int streaming_batch_size;
+  int device_buffer_samples;
 };
 
 template <typename T>
@@ -416,7 +416,7 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam<KmeansBatchedInputs
                                raft::make_host_scalar_view<int>(&ref_n_iter));
 
     cuvs::cluster::kmeans::params batched_params = params;
-    batched_params.streaming_batch_size          = testparams.streaming_batch_size;
+    batched_params.device_buffer_samples         = testparams.device_buffer_samples;
 
     std::optional<raft::host_vector_view<const T, int64_t>> h_sw = std::nullopt;
     auto h_sample_weight = raft::make_host_vector<T, int64_t>(testparams.weighted ? n_samples : 0);
@@ -496,15 +496,15 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam<KmeansBatchedInputs
     auto stream    = raft::resource::get_cuda_stream(handle);
 
     cuvs::cluster::kmeans::params p;
-    p.n_clusters           = n_clusters;
-    p.tol                  = testparams.tol;
-    p.n_init               = 1;
-    p.init                 = cuvs::cluster::kmeans::params::KMeansPlusPlus;
-    p.max_iter             = 20;
-    p.rng_state.seed       = 1;
-    p.oversampling_factor  = 0;
-    p.streaming_batch_size = testparams.streaming_batch_size;
-    p.init_size            = init_size_value;
+    p.n_clusters            = n_clusters;
+    p.tol                   = testparams.tol;
+    p.n_init                = 1;
+    p.init                  = cuvs::cluster::kmeans::params::KMeansPlusPlus;
+    p.max_iter              = 20;
+    p.rng_state.seed        = 1;
+    p.oversampling_factor   = 0;
+    p.device_buffer_samples = testparams.device_buffer_samples;
+    p.init_size             = init_size_value;
 
     auto d_centroids_buf = raft::make_device_matrix<T, int64_t>(handle, n_clusters, n_features);
     T inertia            = 0;