From 7df2d188a0a8924140f689aface37f294868b968 Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Tue, 11 Aug 2026 21:28:39 +0000
Subject: [PATCH 1/2] update batching

---
 cpp/src/cluster/detail/kmeans.cuh | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh
index e3ffb4a439..8225d85479 100644
--- a/cpp/src/cluster/detail/kmeans.cuh
+++ b/cpp/src/cluster/detail/kmeans.cuh
@@ -24,6 +24,7 @@
 #include <raft/core/pinned_mdarray.hpp>
 #include <raft/core/pinned_mdspan.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/linalg/map.cuh>
@@ -696,15 +697,29 @@ void kmeans_fit(
 
   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, device_buffer_samples, stream);
+  auto large_workspace_mr = raft::resource::get_large_workspace_resource_ref(handle);
+  auto data_batches =
+    cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
+      handle,
+      X.data_handle(),
+      n_samples,
+      n_features,
+      device_buffer_samples,
+      stream,
+      large_workspace_mr);
   // 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}, device_buffer_samples, stream);
+        handle,
+        weight_ptr,
+        n_samples,
+        IndexT{1},
+        device_buffer_samples,
+        stream,
+        large_workspace_mr);
     } else {
       raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1});
     }

From 3b432ee2d48099cc387d3b6f99f15de0fc63951d Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Tue, 11 Aug 2026 21:39:54 +0000
Subject: [PATCH 2/2] update workspace conditional

---
 cpp/src/cluster/detail/kmeans.cuh | 32 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh
index 8225d85479..0a429d0528 100644
--- a/cpp/src/cluster/detail/kmeans.cuh
+++ b/cpp/src/cluster/detail/kmeans.cuh
@@ -697,29 +697,27 @@ void kmeans_fit(
 
   rmm::device_uvector<char> batch_workspace(device_buffer_samples, stream);
 
-  auto large_workspace_mr = raft::resource::get_large_workspace_resource_ref(handle);
-  auto data_batches =
-    cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
-      handle,
-      X.data_handle(),
-      n_samples,
-      n_features,
-      device_buffer_samples,
-      stream,
-      large_workspace_mr);
+  auto batch_memory = raft::resource::get_workspace_resource_ref(handle);
+  if constexpr (!data_on_device) {
+    size_t batch_staging_bytes =
+      static_cast<size_t>(device_buffer_samples) * static_cast<size_t>(n_features) * sizeof(DataT);
+    if (weight_ptr != nullptr) {
+      batch_staging_bytes += static_cast<size_t>(device_buffer_samples) * sizeof(DataT);
+    }
+    if (batch_staging_bytes > raft::resource::get_workspace_free_bytes(handle)) {
+      batch_memory = raft::resource::get_large_workspace_resource_ref(handle);
+    }
+  }
+
+  auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
+    handle, X.data_handle(), n_samples, n_features, device_buffer_samples, stream, batch_memory);
   // 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},
-        device_buffer_samples,
-        stream,
-        large_workspace_mr);
+        handle, weight_ptr, n_samples, IndexT{1}, device_buffer_samples, stream, batch_memory);
     } else {
       raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1});
     }