From 1e9160a1ddea57cea0eed10f2be8a8374efd18fa Mon Sep 17 00:00:00 2001 From: vic Date: Tue, 28 Jul 2026 12:03:55 +0200 Subject: [PATCH 1/3] Use batched unfused 1-NN for KMeans on Blackwell --- .../detail/minClusterDistanceCompute.cu | 104 ++++++++++++------ 1 file changed, 72 insertions(+), 32 deletions(-) diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index b15119599e..ee3cc3cdfd 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -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 */ @@ -27,15 +27,15 @@ void minClusterAndDistanceCompute( int batch_centroids, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - bool is_fused = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - if (is_fused) { + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto n_samples = X.extent(0); + auto n_features = X.extent(1); + auto n_clusters = centroids.extent(0); + const bool can_use_fused_path = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + + if (can_use_fused_path) { L2NormBuf_OR_DistBuf.resize(n_clusters, stream); auto centroidsNorm = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); @@ -51,10 +51,10 @@ void minClusterAndDistanceCompute( raft::KeyValuePair initial_value(0, std::numeric_limits::max()); raft::matrix::fill(handle, minClusterAndDistance, initial_value); - bool should_use_fused = + const bool use_fused_path = use_fused(handle, n_samples, n_clusters, n_features); - if (should_use_fused) { + if (use_fused_path) { workspace.resize((sizeof(int)) * n_samples, stream); cuvs::distance::fusedDistanceNNMinReduce, IndexT>( @@ -74,26 +74,66 @@ void minClusterAndDistanceCompute( 0.0f, stream); } else { - workspace.resize(sizeof(DataT) * n_samples * n_clusters, stream); - - cuvs::distance:: - unfusedDistanceNNMinReduce, IndexT>( - handle, - minClusterAndDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); + auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); + auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); + + // The unfused reduction indexes its distance matrix with IndexT. + dataBatchSize = + std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); + + workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); + + using KeyValueT = raft::KeyValuePair; + const bool tileCentroids = centroidsBatchSize < n_clusters; + rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, + stream); + + for (IndexT dIdx = 0; dIdx < n_samples;) { + auto ns = std::min(dataBatchSize, n_samples - dIdx); + auto minClusterAndDistanceView = raft::make_device_vector_view( + minClusterAndDistance.data_handle() + dIdx, ns); + + for (IndexT cIdx = 0; cIdx < n_clusters;) { + auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); + auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() + : minClusterAndDistanceView.data_handle(); + + cuvs::distance::unfusedDistanceNNMinReduce( + handle, + batchMin, + X.data_handle() + dIdx * n_features, + centroids.data_handle() + cIdx * n_features, + L2NormX.data_handle() + dIdx, + centroidsNorm.data_handle() + cIdx, + ns, + nc, + n_features, + (void*)workspace.data(), + metric != cuvs::distance::DistanceType::L2Expanded, + tileCentroids, + true, + metric, + 0.0f, + stream); + + if (tileCentroids) { + // Convert tile-local centroid indices and merge the tile minima. + auto batchMinView = + raft::make_device_vector_view(batchMin, ns); + raft::linalg::map( + handle, + minClusterAndDistanceView, + [cIdx] __device__(KeyValueT current, KeyValueT batch) { + batch.key += cIdx; + return batch.value < current.value ? batch : current; + }, + raft::make_const_mdspan(minClusterAndDistanceView), + batchMinView); + } + cIdx += nc; + } + dIdx += ns; + } } } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); From 711b3f6efcc247b6643dcf36c083936421931342 Mon Sep 17 00:00:00 2001 From: vic Date: Tue, 28 Jul 2026 17:00:04 +0200 Subject: [PATCH 2/3] Add testing --- cpp/tests/CMakeLists.txt | 9 +- cpp/tests/cluster/kmeans_predict_batching.cu | 156 +++++++++++++++++++ 2 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 cpp/tests/cluster/kmeans_predict_batching.cu diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c0decb2243..437396b736 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -123,8 +123,13 @@ ConfigureTest( ConfigureTest( NAME CLUSTER_TEST - PATH cluster/kmeans.cu cluster/kmeans_balanced.cu cluster/kmeans_find_k.cu cluster/linkage.cu - cluster/connect_knn.cu cluster/spectral.cu + PATH cluster/kmeans.cu + cluster/kmeans_balanced.cu + cluster/kmeans_find_k.cu + cluster/kmeans_predict_batching.cu + cluster/linkage.cu + cluster/connect_knn.cu + cluster/spectral.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/cluster/kmeans_predict_batching.cu b/cpp/tests/cluster/kmeans_predict_batching.cu new file mode 100644 index 0000000000..bcff2c4a92 --- /dev/null +++ b/cpp/tests/cluster/kmeans_predict_batching.cu @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include "../../src/cluster/detail/kmeans_common.cuh" + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::cluster::kmeans { + +struct BatchConfig { + char const* name; + int batch_samples; + int batch_centroids; +}; + +constexpr int test_n_samples = 8; +constexpr int test_n_clusters = 8; +constexpr int test_n_features = 2; + +template +std::size_t run_predict_with_batching(raft::resources const& handle, BatchConfig config) +{ + SCOPED_TRACE(config.name); + + using DataT = float; + + constexpr IndexT n_samples = test_n_samples; + constexpr IndexT n_clusters = test_n_clusters; + constexpr IndexT n_features = test_n_features; + + // The nearest centroids span all three centroid tiles. The batched configurations also have + // partial final tiles, exercising the tile-local index adjustment and the minima merge. + constexpr std::array h_x{0.0f, + 0.0f, + 0.5f, + 0.5f, + 3.0f, + 3.0f, + 3.25f, + 3.25f, + 6.0f, + 6.0f, + 5.75f, + 5.75f, + 2.0f, + 2.0f, + 4.0f, + 4.0f}; + constexpr std::array h_centroids{0.0f, + 0.0f, + 100.0f, + 100.0f, + 200.0f, + 200.0f, + 3.0f, + 3.0f, + 400.0f, + 400.0f, + 500.0f, + 500.0f, + 6.0f, + 6.0f, + 700.0f, + 700.0f}; + constexpr std::array expected_clusters{0, 0, 3, 3, 6, 6, 3, 3}; + constexpr DataT expected_inertia = 4.75f; + + auto stream = raft::resource::get_cuda_stream(handle); + auto x = raft::make_device_matrix(handle, n_samples, n_features); + auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + auto labels = raft::make_device_vector(handle, n_samples); + + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream); + raft::update_device(centroids.data_handle(), h_centroids.data(), h_centroids.size(), stream); + + params kmeans_params; + kmeans_params.n_clusters = n_clusters; + kmeans_params.batch_samples = config.batch_samples; + kmeans_params.batch_centroids = config.batch_centroids; + + DataT inertia = 0; + std::size_t total_device_bytes = 0; + { + raft::memory_stats_resources tracked_handle{handle}; + predict(tracked_handle, + kmeans_params, + raft::make_const_mdspan(x.view()), + std::optional>{std::nullopt}, + raft::make_const_mdspan(centroids.view()), + labels.view(), + false, + raft::make_host_scalar_view(&inertia)); + total_device_bytes = tracked_handle.get_bytes_total_allocated().device_global; + } + + std::array h_labels; + raft::update_host(h_labels.data(), labels.data_handle(), h_labels.size(), stream); + raft::resource::sync_stream(handle); + + EXPECT_NEAR(inertia, expected_inertia, 1e-4f); + for (std::size_t i = 0; i < h_labels.size(); ++i) { + EXPECT_EQ(h_labels[i], expected_clusters[i]) << "sample " << i; + } + + return total_device_bytes; +} + +TEST(KMeansPredict, BatchParametersPreserveResultsAndReduceUnfusedAllocations) +{ + raft::resources handle; + constexpr std::array batch_configs{{ + {"no batching", 0, 0}, + {"samples only", 3, 0}, + {"centroids only", 0, 3}, + {"samples and centroids", 3, 3}, + }}; + + auto unbatched_int_bytes = run_predict_with_batching(handle, batch_configs.front()); + auto unbatched_int64_bytes = run_predict_with_batching(handle, batch_configs.front()); + + // predict selects fused or unfused 1-NN according to the architecture heuristic. The batching + // parameters only affect the unfused path, so every GPU checks the results while allocation + // reductions are required only when this problem shape dispatches to unfused 1-NN. + const bool uses_unfused_path = + !detail::use_fused(handle, test_n_samples, test_n_clusters, test_n_features); + + for (std::size_t i = 1; i < batch_configs.size(); ++i) { + auto config = batch_configs[i]; + auto int_bytes = run_predict_with_batching(handle, config); + auto int64_bytes = run_predict_with_batching(handle, config); + + if (uses_unfused_path) { + // Before the batching fix every configuration allocated the full + // samples-by-centroids matrix. + EXPECT_LT(int_bytes, unbatched_int_bytes) << config.name; + EXPECT_LT(int64_bytes, unbatched_int64_bytes) << config.name; + } + } +} + +} // namespace cuvs::cluster::kmeans From 67c0a9a049c1d2482c9c6e4472bc1316c5f539e6 Mon Sep 17 00:00:00 2001 From: vic Date: Wed, 29 Jul 2026 10:03:17 +0200 Subject: [PATCH 3/3] update testing comment --- cpp/tests/cluster/kmeans_predict_batching.cu | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/tests/cluster/kmeans_predict_batching.cu b/cpp/tests/cluster/kmeans_predict_batching.cu index bcff2c4a92..e22cbc9dcd 100644 --- a/cpp/tests/cluster/kmeans_predict_batching.cu +++ b/cpp/tests/cluster/kmeans_predict_batching.cu @@ -145,8 +145,7 @@ TEST(KMeansPredict, BatchParametersPreserveResultsAndReduceUnfusedAllocations) auto int64_bytes = run_predict_with_batching(handle, config); if (uses_unfused_path) { - // Before the batching fix every configuration allocated the full - // samples-by-centroids matrix. + // Verify that batching uses less memory than the unbatched path. EXPECT_LT(int_bytes, unbatched_int_bytes) << config.name; EXPECT_LT(int64_bytes, unbatched_int64_bytes) << config.name; }