[FEA] Skip Device Concat for KmeansMG

cpp/include/cuml/cluster/kmeans.hpp

@@ -23,7 +23,7 @@ namespace kmeans {
  manages the CUDA resources.
  * @param[in]     params        Parameters for KMeans model.
  * @param[in]     X             Training instances to cluster, in row-major
- * format. May live on the device or on the host.
+ * format. May or may not be device accessible.
  * @param[in]     n_samples     Number of samples in the input X.
  * @param[in]     n_features    Number of features or the dimensions of each
  * sample.
@@ -33,7 +33,7 @@ namespace kmeans {
  centroids  as the initial cluster centers
  *                              [out] Otherwise, generated centroids from the
  kmeans algorithm is stored at the address pointed by 'centroids'. `centroids`
- * must always live on the device.
+ * must always be device accessible.
  * @param[out]    inertia       Sum of squared distances of samples to their
  closest cluster center.
  * @param[out]    n_iter        Number of iterations run.
@@ -78,6 +78,57 @@ void fit(const raft::handle_t& handle,
          double& inertia,
          int64_t& n_iter);
 
+/**
+ * @brief Multi-GPU / out-of-core k-means fit over multiple local data partitions.
+ *
+ * Each rank (e.g. each Dask worker) supplies its local training data as an
+ * array of `n_parts` partitions.
+ * All partitions on a given rank must share the same residency. The
+ * distributed reduction across ranks is performed via the NCCL communicator
+ * that must be initialized on `handle`.
+ *
+ * @param[in]     handle              cuML handle with NCCL comms initialized.
+ * @param[in]     params              Parameters for the KMeans model. For
+ *                                    host-resident partitions the
+ *                                    host-to-device batch size is read from
+ *                                    `params.device_buffer_samples`.
+ * @param[in]     X_parts             Array of `n_parts` pointers to the local
+ *                                    row-major partitions (all host- or all
+ *                                    device-resident). Partition `i` has shape
+ *                                    [`n_samples_parts[i]`, `n_features`].
+ * @param[in]     n_samples_parts     Array of `n_parts` per-partition row counts.
+ * @param[in]     n_parts             Number of local partitions on this rank.
+ * @param[in]     n_features          Number of features (shared by all partitions).
+ * @param[in]     sample_weight_parts Optional array of `n_parts` pointers to the
+ *                                    per-partition weight vectors (matching the
+ *                                    residency of `X_parts`), or `nullptr` for
+ *                                    uniform weights.
+ * @param[inout]  centroids           Device matrix [n_clusters x n_features].
+ * @param[out]    inertia             Sum of squared distances to closest center.
+ * @param[out]    n_iter              Number of iterations run.
+ */
+void fit(const raft::handle_t& handle,
+         const KMeansParams& params,
+         const float* const* X_parts,
+         const int64_t* n_samples_parts,
+         int64_t n_parts,
+         int64_t n_features,
+         const float* const* sample_weight_parts,
+         float* centroids,
+         float& inertia,
+         int64_t& n_iter);
+
+void fit(const raft::handle_t& handle,
+         const KMeansParams& params,
+         const double* const* X_parts,
+         const int64_t* n_samples_parts,
+         int64_t n_parts,
+         int64_t n_features,
+         const double* const* sample_weight_parts,
+         double* centroids,
+         double& inertia,
+         int64_t& n_iter);
+
 /**
  * @brief Predict the closest cluster each sample in X belongs to.
  *

cpp/src/kmeans/kmeans_fit.cu

@@ -15,10 +15,89 @@
 #include <kmeans/kmeans_params.hpp>
 
 #include <optional>
+#include <vector>
 
 namespace ML {
 namespace kmeans {
 
+template <typename value_t>
+void fit_impl_device_parts(const raft::handle_t& handle,
+                           const KMeansParams& params,
+                           const value_t* const* X_parts,
+                           const int64_t* n_samples_parts,
+                           int64_t n_parts,
+                           int64_t n_features,
+                           const value_t* const* sample_weight_parts,
+                           value_t* centroids,
+                           value_t& inertia,
+                           int64_t& n_iter)
+{
+  std::vector<raft::device_matrix_view<const value_t, int64_t>> X_views;
+  X_views.reserve(n_parts);
+  for (int64_t i = 0; i < n_parts; ++i) {
+    X_views.push_back(raft::make_device_matrix_view<const value_t, int64_t>(
+      X_parts[i], n_samples_parts[i], n_features));
+  }
+
+  std::optional<std::vector<raft::device_vector_view<const value_t, int64_t>>> sw = std::nullopt;
+  if (sample_weight_parts != nullptr) {
+    std::vector<raft::device_vector_view<const value_t, int64_t>> sw_views;
+    sw_views.reserve(n_parts);
+    for (int64_t i = 0; i < n_parts; ++i) {
+      sw_views.push_back(raft::make_device_vector_view<const value_t, int64_t>(
+        sample_weight_parts[i], n_samples_parts[i]));
+    }
+    sw = std::make_optional(std::move(sw_views));
+  }
+
+  auto centroids_view =
+    raft::make_device_matrix_view<value_t, int64_t>(centroids, params.n_clusters, n_features);
+  auto inertia_view = raft::make_host_scalar_view<value_t>(&inertia);
+  auto n_iter_view  = raft::make_host_scalar_view<int64_t>(&n_iter);
+
+  cuvs::cluster::kmeans::fit(
+    handle, to_cuvs(params), X_views, sw, centroids_view, inertia_view, n_iter_view);
+}
+
+template <typename value_t>
+void fit_impl_host_parts(const raft::handle_t& handle,
+                         const KMeansParams& params,
+                         const value_t* const* X_parts,
+                         const int64_t* n_samples_parts,
+                         int64_t n_parts,
+                         int64_t n_features,
+                         const value_t* const* sample_weight_parts,
+                         value_t* centroids,
+                         value_t& inertia,
+                         int64_t& n_iter)
+{
+  std::vector<raft::host_matrix_view<const value_t, int64_t>> X_views;
+  X_views.reserve(n_parts);
+  for (int64_t i = 0; i < n_parts; ++i) {
+    X_views.push_back(raft::make_host_matrix_view<const value_t, int64_t>(
+      X_parts[i], n_samples_parts[i], n_features));
+  }
+
+  std::optional<std::vector<raft::host_vector_view<const value_t, int64_t>>> sw = std::nullopt;
+  if (sample_weight_parts != nullptr) {
+    std::vector<raft::host_vector_view<const value_t, int64_t>> sw_views;
+    sw_views.reserve(n_parts);
+    for (int64_t i = 0; i < n_parts; ++i) {
+      sw_views.push_back(raft::make_host_vector_view<const value_t, int64_t>(sample_weight_parts[i],
+                                                                             n_samples_parts[i]));
+    }
+    sw = std::make_optional(std::move(sw_views));
+  }
+
+  auto centroids_view =
+    raft::make_device_matrix_view<value_t, int64_t>(centroids, params.n_clusters, n_features);
+  auto inertia_view = raft::make_host_scalar_view<value_t>(&inertia);
+  auto n_iter_view  = raft::make_host_scalar_view<int64_t>(&n_iter);
+
+  cuvs::cluster::kmeans::fit(
+    handle, to_cuvs(params), X_views, sw, centroids_view, inertia_view, n_iter_view);
+}
+
 template <typename value_t, typename idx_t>
 void fit_impl_host(const raft::handle_t& handle,
                    const KMeansParams& params,
@@ -151,5 +230,92 @@ void fit(const raft::handle_t& handle,
   }
 }
 
+// Detect partition residency from the first non-empty partition: an empty
+// partition may carry a null data pointer that `is_device_or_managed_type`
+// cannot classify, so skip past empties. An all-empty rank has no local data;
+// default to the host path.
+template <typename value_t>
+static bool parts_on_device(const value_t* const* X_parts,
+                            const int64_t* n_samples_parts,
+                            int64_t n_parts)
+{
+  for (int64_t i = 0; i < n_parts; ++i) {
+    if (n_samples_parts[i] > 0) { return ML::is_device_or_managed_type(X_parts[i]); }
+  }
+  return false;
+}
+
+void fit(const raft::handle_t& handle,
+         const KMeansParams& params,
+         const float* const* X_parts,
+         const int64_t* n_samples_parts,
+         int64_t n_parts,
+         int64_t n_features,
+         const float* const* sample_weight_parts,
+         float* centroids,
+         float& inertia,
+         int64_t& n_iter)
+{
+  if (parts_on_device(X_parts, n_samples_parts, n_parts)) {
+    fit_impl_device_parts(handle,
+                          params,
+                          X_parts,
+                          n_samples_parts,
+                          n_parts,
+                          n_features,
+                          sample_weight_parts,
+                          centroids,
+                          inertia,
+                          n_iter);
+  } else {
+    fit_impl_host_parts(handle,
+                        params,
+                        X_parts,
+                        n_samples_parts,
+                        n_parts,
+                        n_features,
+                        sample_weight_parts,
+                        centroids,
+                        inertia,
+                        n_iter);
+  }
+}
+
+void fit(const raft::handle_t& handle,
+         const KMeansParams& params,
+         const double* const* X_parts,
+         const int64_t* n_samples_parts,
+         int64_t n_parts,
+         int64_t n_features,
+         const double* const* sample_weight_parts,
+         double* centroids,
+         double& inertia,
+         int64_t& n_iter)
+{
+  if (parts_on_device(X_parts, n_samples_parts, n_parts)) {
+    fit_impl_device_parts(handle,
+                          params,
+                          X_parts,
+                          n_samples_parts,
+                          n_parts,
+                          n_features,
+                          sample_weight_parts,
+                          centroids,
+                          inertia,
+                          n_iter);
+  } else {
+    fit_impl_host_parts(handle,
+                        params,
+                        X_parts,
+                        n_samples_parts,
+                        n_parts,
+                        n_features,
+                        sample_weight_parts,
+                        centroids,
+                        inertia,
+                        n_iter);
+  }
+}
+
 };  // end namespace kmeans
 };  // end namespace ML

python/cuml/cuml/cluster/cpp/kmeans.pxd

@@ -113,6 +113,28 @@ cdef extern from "cuml/cluster/kmeans.hpp" namespace "ML::kmeans" nogil:
                   double &inertia,
                   int64_t &n_iter) except +
 
+    cdef void fit(handle_t& handle,
+                  KMeansParams& params,
+                  float **X_parts,
+                  const int64_t *n_samples_parts,
+                  int64_t n_parts,
+                  int64_t n_features,
+                  float **sample_weight_parts,
+                  float *centroids,
+                  float &inertia,
+                  int64_t &n_iter) except +
+
+    cdef void fit(handle_t& handle,
+                  KMeansParams& params,
+                  double **X_parts,
+                  const int64_t *n_samples_parts,
+                  int64_t n_parts,
+                  int64_t n_features,
+                  double **sample_weight_parts,
+                  double *centroids,
+                  double &inertia,
+                  int64_t &n_iter) except +
+
     cdef void predict(handle_t& handle,
                       KMeansParams& params,
                       const float *centroids,