diff --git a/c/include/cuvs/cluster/kmeans.h b/c/include/cuvs/cluster/kmeans.h index 9a3882cb4c..b57a678336 100644 --- a/c/include/cuvs/cluster/kmeans.h +++ b/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,10 +188,10 @@ 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. @@ -199,6 +199,14 @@ struct cuvsKMeansParams { * 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; diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 8e46764ce4..0fdb35894e 100644 --- a/c/src/cluster/kmeans.cpp +++ b/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 -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(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(params.metric); + kmeans_params.init = static_cast(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(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(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}; }); } diff --git a/c/src/cluster/mg_kmeans.cpp b/c/src/cluster/mg_kmeans.cpp index 1f1b51ff00..977a12a244 100644 --- a/c/src/cluster/mg_kmeans.cpp +++ b/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; } diff --git a/c/tests/cluster/kmeans_mg_c.cu b/c/tests/cluster/kmeans_mg_c.cu index 40322e8057..4d8cde6968 100644 --- a/c/tests/cluster/kmeans_mg_c.cu +++ b/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(¶ms), 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( diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index 74ad6d09d4..7ce160f113 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/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]. diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index b1608407bc..921d97abe5 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/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(pams.streaming_batch_size); - if (streaming_batch_size <= 0 || streaming_batch_size > static_cast(n_samples)) { - streaming_batch_size = static_cast(n_samples); + IndexT device_buffer_batch_size = static_cast(pams.device_buffer_batch_size); + if (device_buffer_batch_size <= 0 || device_buffer_batch_size > static_cast(n_samples)) { + device_buffer_batch_size = static_cast(n_samples); } constexpr bool data_on_device = raft::is_device_mdspan_v; @@ -606,13 +606,13 @@ void kmeans_fit( rmm::device_uvector local_workspace(0, stream); rmm::device_uvector& ws = workspace.has_value() ? workspace->get() : local_workspace; - if (data_on_device && streaming_batch_size != static_cast(n_samples)) { + if (data_on_device && device_buffer_batch_size != static_cast(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(streaming_batch_size), + static_cast(device_buffer_batch_size), static_cast(n_samples)); - streaming_batch_size = static_cast(n_samples); + device_buffer_batch_size = static_cast(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, IndexT>( - handle, streaming_batch_size); - auto L2NormBatch = raft::make_device_vector(handle, streaming_batch_size); - auto batch_weights_buf = raft::make_device_vector(handle, streaming_batch_size); + handle, device_buffer_batch_size); + auto L2NormBatch = raft::make_device_vector(handle, device_buffer_batch_size); + auto batch_weights_buf = + raft::make_device_vector(handle, device_buffer_batch_size); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); auto centroid_sums = raft::make_device_matrix(handle, n_clusters, n_features); auto weight_per_cluster = raft::make_device_vector(handle, n_clusters); auto clustering_cost = raft::make_device_scalar(handle, DataT{0}); - rmm::device_uvector batch_workspace(streaming_batch_size, stream); + rmm::device_uvector batch_workspace(device_buffer_batch_size, stream); auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - 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> weight_batches; if constexpr (!data_on_device) { if (weight_ptr != nullptr) { weight_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - 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(n_samples), static_cast(n_features), n_clusters, - static_cast(streaming_batch_size)); + static_cast(device_buffer_batch_size)); bool need_compute_norms = metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded; diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index c83b3fd22f..9d78bd6c80 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -37,6 +37,7 @@ #include #include +#include #include #include @@ -44,6 +45,7 @@ #include #include #include +#include #include #include #include @@ -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 data_copy_stream_owner; + rmm::cuda_stream_view data_copy_stream{stream}; + bool enable_data_prefetch = false; + if constexpr (!data_on_device) { + if (params.device_buffer_prefetch) { + data_copy_stream_owner = + std::make_unique(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(params.n_clusters); auto metric = params.metric; @@ -188,18 +201,18 @@ void mnmg_fit( static_cast(n_features), static_cast(n_clusters)); - IndexT streaming_batch_size = static_cast(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_batch_size = static_cast(params.device_buffer_batch_size); + if (device_buffer_batch_size <= 0 || device_buffer_batch_size > max_part_rows) { + device_buffer_batch_size = std::max(max_part_rows, IndexT{1}); } - if (data_on_device && streaming_batch_size < max_part_rows) { + if (data_on_device && device_buffer_batch_size < max_part_rows) { RAFT_LOG_WARN( - "MNMG KMeans: streaming_batch_size (%zu) ignored when partitions reside on device; using " + "MNMG KMeans: device_buffer_batch_size (%zu) ignored when partitions reside on device; using " "max partition size (%zu)", - static_cast(streaming_batch_size), + static_cast(device_buffer_batch_size), static_cast(max_part_rows)); - streaming_batch_size = max_part_rows; + device_buffer_batch_size = max_part_rows; } auto rank_centroids_arr = @@ -212,7 +225,7 @@ void mnmg_fit( auto clustering_cost = raft::make_device_vector(dev_res, 1); auto batch_clustering_cost = raft::make_device_vector(dev_res, 1); auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); - IndexT alloc_batch_size = streaming_batch_size; + IndexT alloc_batch_size = device_buffer_batch_size; auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); auto minClusterAndDistance = raft::make_device_vector, IndexT>(dev_res, alloc_batch_size); @@ -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 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(device_buffer_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(); @@ -363,7 +394,7 @@ void mnmg_fit( init_centroids_for_mg_batched(dev_res, iter_params, - streaming_batch_size, + device_buffer_batch_size, X_parts, n_features, input_centroids_const, @@ -374,6 +405,15 @@ void mnmg_fit( rank, comms); + if (persist_data_batches) { + persistent_data_batches.emplace(dev_res, + X_parts.front(), + static_cast(device_buffer_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(X_part.extent(0)); if (part_rows == 0) { continue; } - data_batch_iterator_t data_batches(dev_res, - X_part, - static_cast(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(data_batch.size()); auto batch_offset = static_cast(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(X_part.extent(0)); if (part_rows == 0) { continue; } - data_batch_iterator_t data_batches(dev_res, - X_part, - static_cast(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(data_batch.size()); auto batch_offset = static_cast(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); diff --git a/cpp/src/cluster/detail/kmeans_mg_batched_init.cuh b/cpp/src/cluster/detail/kmeans_mg_batched_init.cuh index b6cb021d2c..49837f9ee9 100644 --- a/cpp/src/cluster/detail/kmeans_mg_batched_init.cuh +++ b/cpp/src/cluster/detail/kmeans_mg_batched_init.cuh @@ -247,7 +247,7 @@ template void init_centroids_for_mg_batched( raft::resources const& handle, const cuvs::cluster::kmeans::params& params, - IndexT /*streaming_batch_size*/, + IndexT /*device_buffer_batch_size*/, const std::vector< raft::mdspan, raft::row_major, Accessor>>& X_parts, IndexT n_features, diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 0b5d1b8fc9..70352d789a 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/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_batch_size; }; template @@ -416,7 +416,7 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(&ref_n_iter)); cuvs::cluster::kmeans::params batched_params = params; - batched_params.streaming_batch_size = testparams.streaming_batch_size; + batched_params.device_buffer_batch_size = testparams.device_buffer_batch_size; std::optional> h_sw = std::nullopt; auto h_sample_weight = raft::make_host_vector(testparams.weighted ? n_samples : 0); @@ -496,15 +496,15 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(handle, n_clusters, n_features); T inertia = 0; diff --git a/cpp/tests/cluster/kmeans_mg.cu b/cpp/tests/cluster/kmeans_mg.cu index 0830dec040..fcda9708b7 100644 --- a/cpp/tests/cluster/kmeans_mg.cu +++ b/cpp/tests/cluster/kmeans_mg.cu @@ -38,10 +38,11 @@ struct KmeansSNMGInputs { int n_clusters; T tol; kmeans_weight_mode weight_mode; - int streaming_batch_size; + int device_buffer_batch_size; int n_init; cuvs::cluster::kmeans::params::InitMethod init = cuvs::cluster::kmeans::params::Array; int max_iter = 20; + bool device_buffer_prefetch = false; }; template @@ -105,13 +106,14 @@ class KmeansSNMGTest : public ::testing::TestWithParam> { } 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.device_buffer_batch_size = testparams_.device_buffer_batch_size; + snmg_params.device_buffer_prefetch = testparams_.device_buffer_prefetch; T snmg_inertia = T{0}; int64_t snmg_n_iter = 0; @@ -277,10 +279,32 @@ class KmeansSNMGTest : public ::testing::TestWithParam> { // SNMG float test inputs // ============================================================================ const std::vector> snmg_inputsf = { - // n_row, n_col, n_clusters, tol, weight_mode, streaming_batch_size, n_init[, init] + // n_row, n_col, n_clusters, tol, weight_mode, device_buffer_batch_size, n_init[, init] {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}, diff --git a/cpp/tests/cluster/kmeans_mnmg.cu b/cpp/tests/cluster/kmeans_mnmg.cu index 39feeef1b6..af4a276f02 100644 --- a/cpp/tests/cluster/kmeans_mnmg.cu +++ b/cpp/tests/cluster/kmeans_mnmg.cu @@ -278,7 +278,7 @@ struct KmeansMGNcclInputs { int n_clusters; T tol; kmeans_weight_mode weight_mode; - int64_t streaming_batch_size; + int64_t device_buffer_batch_size; int n_init; int partitions_per_rank; cuvs::cluster::kmeans::params::InitMethod init = cuvs::cluster::kmeans::params::Array; @@ -359,13 +359,13 @@ class KmeansMGNcclTest : public ::testing::TestWithParam> } cuvs::cluster::kmeans::params kp; - kp.n_clusters = n_clusters; - kp.tol = testparams_.tol; - kp.max_iter = testparams_.max_iter; - kp.n_init = testparams_.n_init; - kp.rng_state.seed = 42; - kp.init = testparams_.init; - kp.streaming_batch_size = testparams_.streaming_batch_size; + kp.n_clusters = n_clusters; + kp.tol = testparams_.tol; + kp.max_iter = testparams_.max_iter; + kp.n_init = testparams_.n_init; + kp.rng_state.seed = 42; + kp.init = testparams_.init; + kp.device_buffer_batch_size = testparams_.device_buffer_batch_size; std::vector h_mg_centroids; T mg_inertia = T{0}; @@ -414,13 +414,13 @@ class KmeansMGNcclTest : public ::testing::TestWithParam> } cuvs::cluster::kmeans::params skp; - skp.n_clusters = n_clusters; - skp.tol = testparams_.tol; - skp.max_iter = testparams_.max_iter; - skp.n_init = testparams_.n_init; - skp.rng_state.seed = 42; - skp.init = testparams_.init; - skp.streaming_batch_size = testparams_.streaming_batch_size; + skp.n_clusters = n_clusters; + skp.tol = testparams_.tol; + skp.max_iter = testparams_.max_iter; + skp.n_init = testparams_.n_init; + skp.rng_state.seed = 42; + skp.init = testparams_.init; + skp.device_buffer_batch_size = testparams_.device_buffer_batch_size; T sg_inertia = T{0}; int sg_n_iter = 0; @@ -521,7 +521,7 @@ class KmeansMGNcclTest : public ::testing::TestWithParam> // NCCL float test inputs // ============================================================================ const std::vector> mg_nccl_inputsf = { - // n_row, n_col, n_clusters, tol, weight_mode, streaming_batch_size, n_init, + // n_row, n_col, n_clusters, tol, weight_mode, device_buffer_batch_size, n_init, // partitions_per_rank[, init[, max_iter]] {1000, 32, 5, 0.0001f, kmeans_weight_mode::none, 1000, 1, 1}, {1000, 32, 5, 0.0001f, kmeans_weight_mode::none, 1000, 1, 2}, @@ -706,14 +706,14 @@ class KmeansMGOversamplingTest : public ::testing::Test { double oversampling_factor = 1.0) { cuvs::cluster::kmeans::params kp; - kp.n_clusters = n_clusters; - kp.tol = T(1e-4); - kp.max_iter = 30; - kp.n_init = 1; - kp.rng_state.seed = 42; - kp.init = cuvs::cluster::kmeans::params::KMeansPlusPlus; - kp.streaming_batch_size = n_samples; - kp.oversampling_factor = oversampling_factor; + kp.n_clusters = n_clusters; + kp.tol = T(1e-4); + kp.max_iter = 30; + kp.n_init = 1; + kp.rng_state.seed = 42; + kp.init = cuvs::cluster::kmeans::params::KMeansPlusPlus; + kp.device_buffer_batch_size = n_samples; + kp.oversampling_factor = oversampling_factor; std::vector h_centroids; const int actual_threads = run_mg_fit_omp(device_ids_, diff --git a/python/cuvs/cuvs/cluster/kmeans/kmeans.pxd b/python/cuvs/cuvs/cluster/kmeans/kmeans.pxd index 5e6cbfe9ca..6b43d59b6b 100644 --- a/python/cuvs/cuvs/cluster/kmeans/kmeans.pxd +++ b/python/cuvs/cuvs/cluster/kmeans/kmeans.pxd @@ -55,8 +55,9 @@ cdef extern from "cuvs/cluster/kmeans.h" nogil: int batch_centroids, bool hierarchical, int hierarchical_n_iters, - int64_t streaming_batch_size, - int64_t init_size + int64_t device_buffer_batch_size, + int64_t init_size, + bool device_buffer_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 _device_buffer_prefetch diff --git a/python/cuvs/cuvs/cluster/kmeans/kmeans.pyx b/python/cuvs/cuvs/cluster/kmeans/kmeans.pyx index 2a2859a6b2..b0950012c0 100644 --- a/python/cuvs/cuvs/cluster/kmeans/kmeans.pyx +++ b/python/cuvs/cuvs/cluster/kmeans/kmeans.pyx @@ -85,15 +85,22 @@ cdef class KMeansParams: Number of samples to draw for KMeansPlusPlus initialization with host (out-of-core) data. When set to 0, uses the heuristic min(3 * n_clusters, n_samples). Default: 0. - streaming_batch_size : int - Number of samples to process per GPU batch when fitting with host - (numpy) data. When set to 0, defaults to n_samples (process all - at once). Only used by the batched (host-data) code path. Reducing - streaming_batch_size can help reduce GPU memory pressure but increases + device_buffer_batch_size : int + Number of host-resident samples staged in each GPU batch. When set + to 0, defaults to n_samples (process all at once). Only used by the + host-data code path. Reducing + device_buffer_batch_size can help reduce GPU memory pressure but increases overhead as the number of times centroid adjustments are computed increases. Default: 0 (process all data at once). + device_buffer_prefetch : bool + Prefetch the next device buffer asynchronously while the current buffer + is processed. This overlaps transfers from host to device with + computation using a second device buffer per rank. Ignored by other + KMeans paths. + + Default: False. hierarchical : bool Whether to use hierarchical (balanced) kmeans or not hierarchical_n_iters : int @@ -101,6 +108,7 @@ cdef class KMeansParams: """ def __cinit__(self): + self._device_buffer_prefetch = False cuvsKMeansParamsCreate(&self.params) def __dealloc__(self): @@ -118,7 +126,8 @@ cdef class KMeansParams: batch_centroids=None, inertia_check=None, init_size=None, - streaming_batch_size=None, + device_buffer_batch_size=None, + device_buffer_prefetch=None, hierarchical=None, hierarchical_n_iters=None): if metric is not None: @@ -148,8 +157,10 @@ cdef class KMeansParams: ) if init_size is not None: self.params.init_size = init_size - if streaming_batch_size is not None: - self.params.streaming_batch_size = streaming_batch_size + if device_buffer_batch_size is not None: + self.params.streaming_batch_size = device_buffer_batch_size + if device_buffer_prefetch is not None: + self._device_buffer_prefetch = device_buffer_prefetch if hierarchical is not None: self.params.hierarchical = hierarchical if hierarchical_n_iters is not None: @@ -199,9 +210,13 @@ cdef class KMeansParams: return self.params.init_size @property - def streaming_batch_size(self): + def device_buffer_batch_size(self): return self.params.streaming_batch_size + @property + def device_buffer_prefetch(self): + return self._device_buffer_prefetch + @property def hierarchical(self): return self.params.hierarchical @@ -225,16 +240,16 @@ def fit( When X is a device array (CUDA array interface), standard on-device k-means is used. When X is a host array (numpy ndarray or ``__array_interface__``), data is streamed to the GPU in batches - controlled by ``params.streaming_batch_size``. For large host datasets, consider - reducing ``streaming_batch_size`` to reduce GPU memory usage. + controlled by ``params.device_buffer_batch_size``. For large host datasets, consider + reducing ``device_buffer_batch_size`` to reduce GPU memory usage. Parameters ---------- params : KMeansParams Parameters to use to fit KMeans model. For host data, - ``params.streaming_batch_size`` controls how many samples are sent to the - GPU per batch. + ``params.device_buffer_batch_size`` controls how many samples are staged + on the GPU at once. X : array-like Training instances, shape (m, k). Accepts both device arrays (cupy / CUDA array interface) and host arrays (numpy). @@ -275,7 +290,7 @@ def fit( >>> import numpy as np >>> X_host = np.random.random((10_000_000, 128)).astype(np.float32) - >>> params = KMeansParams(n_clusters=1000, streaming_batch_size=1_000_000) + >>> params = KMeansParams(n_clusters=1000, device_buffer_batch_size=1_000_000) >>> centroids, inertia, n_iter = fit(params, X_host) """ diff --git a/python/cuvs/cuvs/cluster/mg/kmeans/kmeans.pyx b/python/cuvs/cuvs/cluster/mg/kmeans/kmeans.pyx index ed1223271d..2c681f4f96 100644 --- a/python/cuvs/cuvs/cluster/mg/kmeans/kmeans.pyx +++ b/python/cuvs/cuvs/cluster/mg/kmeans/kmeans.pyx @@ -76,6 +76,17 @@ def fit( ``centroids`` is a host NumPy array containing the computed centroids, ``inertia`` is the final objective value, and ``n_iter`` is the number of iterations run. + + Notes + ----- + For small streaming batches, call ``resources.set_memory_pool(...)`` + before ``fit`` to reduce allocator contention between GPU worker threads. + Memory pools are not enabled automatically because they replace the + process-wide RMM resource on each managed device. + + Set ``params.device_buffer_prefetch=True`` to overlap transfers from host + to device with computation. This allocates a second device batch buffer on every rank; + the default single-buffered path minimizes device-memory usage. """ if params.hierarchical: @@ -151,8 +162,9 @@ def fit( params_v2.batch_centroids = params.params.batch_centroids params_v2.hierarchical = params.params.hierarchical params_v2.hierarchical_n_iters = params.params.hierarchical_n_iters - params_v2.streaming_batch_size = params.params.streaming_batch_size + params_v2.device_buffer_batch_size = params.params.streaming_batch_size params_v2.init_size = params.params.init_size + params_v2.device_buffer_prefetch = params._device_buffer_prefetch with cuda_interruptible(): check_cuvs(cuvsMultiGpuKMeansFit( diff --git a/python/cuvs/cuvs/tests/test_kmeans.py b/python/cuvs/cuvs/tests/test_kmeans.py index 210ae06b80..cbabb4b35a 100644 --- a/python/cuvs/cuvs/tests/test_kmeans.py +++ b/python/cuvs/cuvs/tests/test_kmeans.py @@ -1,4 +1,4 @@ -# 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 # @@ -79,11 +79,11 @@ def test_cluster_cost(n_rows, n_cols, n_clusters, dtype): @pytest.mark.parametrize("n_rows", [1000, 5000]) @pytest.mark.parametrize("n_cols", [10, 100]) @pytest.mark.parametrize("n_clusters", [8, 16]) -@pytest.mark.parametrize("streaming_batch_size", [0, 100, 239, 500]) +@pytest.mark.parametrize("device_buffer_batch_size", [0, 100, 239, 500]) @pytest.mark.parametrize("dtype", [np.float64]) @pytest.mark.parametrize("weighted", [False, True]) def test_fit_host_matches_fit_device( - n_rows, n_cols, n_clusters, streaming_batch_size, dtype, weighted + n_rows, n_cols, n_clusters, device_buffer_batch_size, dtype, weighted ): """ Test that fit() with host (numpy) data produces the same centroids as @@ -122,7 +122,7 @@ def test_fit_host_matches_fit_device( init_method="Array", max_iter=20, tol=1e-10, - streaming_batch_size=streaming_batch_size, + device_buffer_batch_size=device_buffer_batch_size, ) centroids_batched, inertia_batched, _ = fit( params_host, diff --git a/python/cuvs/cuvs/tests/test_mg_kmeans.py b/python/cuvs/cuvs/tests/test_mg_kmeans.py index 228abc5f67..69190887f4 100644 --- a/python/cuvs/cuvs/tests/test_mg_kmeans.py +++ b/python/cuvs/cuvs/tests/test_mg_kmeans.py @@ -93,7 +93,10 @@ def assert_inertia_matches_centroids(out, X, sample_weights): @pytest.mark.parametrize("dtype", [np.float32, np.float64]) @pytest.mark.parametrize("init_method", ["Array", "KMeansPlusPlus", "Random"]) @pytest.mark.parametrize("weighted", [False, True]) -def test_mg_kmeans_fit_options(dtype, init_method, weighted): +@pytest.mark.parametrize("device_buffer_prefetch", [False, True]) +def test_mg_kmeans_fit_options( + dtype, init_method, weighted, device_buffer_prefetch +): n_clusters = 4 X, initial_centroids = make_inputs(dtype, n_clusters=n_clusters) resources = MultiGpuResources() @@ -110,8 +113,10 @@ def test_mg_kmeans_fit_options(dtype, init_method, weighted): tol=1e-10, n_init=3 if init_method == "Random" else 1, init_size=X.shape[0], - streaming_batch_size=37, + device_buffer_batch_size=37, + device_buffer_prefetch=device_buffer_prefetch, ) + assert params.device_buffer_prefetch == device_buffer_prefetch centroids = initial_centroids.copy() if init_method == "Array" else None mg_out = mg_kmeans.fit(