Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions c/include/cuvs/cluster/kmeans.h
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
8 changes: 4 additions & 4 deletions c/src/cluster/kmeans.cpp
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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};
});
}
Expand Down Expand Up @@ -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};
});
}
Expand Down
2 changes: 1 addition & 1 deletion c/src/cluster/mg_kmeans.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
6 changes: 3 additions & 3 deletions c/tests/cluster/kmeans_c.cu
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion c/tests/cluster/kmeans_mg_c.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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>(
Expand Down
16 changes: 8 additions & 8 deletions cpp/include/cuvs/cluster/kmeans.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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_samples, which controls how many
* samples to transfer from host to device per batch when processing out-of-core
* data.
*/
Expand Down Expand Up @@ -154,7 +154,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;
};

/**
Expand Down Expand Up @@ -199,7 +199,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
Expand All @@ -221,7 +221,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;
*
Expand All @@ -245,7 +245,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]
Expand Down Expand Up @@ -1648,8 +1648,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`:
Expand All @@ -1665,7 +1665,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].
Expand Down
30 changes: 15 additions & 15 deletions cpp/src/cluster/detail/kmeans.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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)>;
Expand All @@ -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.
Expand Down Expand Up @@ -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});
}
Expand Down Expand Up @@ -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;
Expand Down
22 changes: 11 additions & 11 deletions cpp/src/cluster/detail/kmeans_mg.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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);
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion cpp/src/cluster/detail/kmeans_mg_batched_init.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
24 changes: 12 additions & 12 deletions cpp/tests/cluster/kmeans.cu
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -339,7 +339,7 @@ struct KmeansBatchedInputs {
int n_clusters;
T tol;
bool weighted;
int streaming_batch_size;
int device_buffer_samples;
};

template <typename T>
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down
Loading
Loading