Use RAFT backend for make_blobs

python/cuml/cuml/datasets/CMakeLists.txt

@@ -1,12 +1,13 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
 
 set(cython_sources "")
 add_module_gpu_default("arima.pyx" ${arima_algo} ${datasets_algo})
+add_module_gpu_default("_blobs.pyx" ${datasets_algo})
 add_module_gpu_default("regression.pyx" ${regression_algo} ${datasets_algo})
 
 rapids_cython_create_modules(

python/cuml/cuml/datasets/_blobs.pyx

@@ -0,0 +1,116 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+
+import cupy as cp
+
+from cuml.internals import get_handle
+
+from libc.stddef cimport size_t
+from libc.stdint cimport int64_t, uint64_t, uintptr_t
+from libcpp cimport bool
+from pylibraft.common.handle cimport handle_t
+
+
+cdef extern from "cuml/datasets/make_blobs.hpp" namespace "ML" nogil:
+    void cpp_make_blobs "ML::Datasets::make_blobs" (
+        const handle_t& handle,
+        float* out,
+        int64_t* labels,
+        int64_t n_rows,
+        int64_t n_cols,
+        int64_t n_clusters,
+        bool row_major,
+        const float* centers,
+        const float* cluster_std,
+        const float cluster_std_scalar,
+        bool shuffle,
+        float center_box_min,
+        float center_box_max,
+        uint64_t seed) except +
+
+    void cpp_make_blobs "ML::Datasets::make_blobs" (
+        const handle_t& handle,
+        double* out,
+        int64_t* labels,
+        int64_t n_rows,
+        int64_t n_cols,
+        int64_t n_clusters,
+        bool row_major,
+        const double* centers,
+        const double* cluster_std,
+        const double cluster_std_scalar,
+        bool shuffle,
+        double center_box_min,
+        double center_box_max,
+        uint64_t seed) except +
+
+
+def make_blobs(
+    n_samples,
+    n_features,
+    n_centers,
+    centers,
+    cluster_std,
+    center_box_min,
+    center_box_max,
+    shuffle,
+    random_state,
+    order,
+    dtype,
+):
+    dtype = cp.dtype(dtype)
+
+    h = get_handle()
+    cdef handle_t* h_ptr = <handle_t*><size_t>h.getHandle()
+
+    X = cp.empty((n_samples, n_features), dtype=dtype, order=order)
+    y = cp.empty(n_samples, dtype=cp.int64)
+
+    cdef uintptr_t x_p = X.data.ptr
+    cdef uintptr_t y_p = y.data.ptr
+    cdef uintptr_t ctr_p = 0
+    if centers is not None:
+        ctr_p = centers.data.ptr
+
+    cdef bool row_c = order == "C"
+
+    if dtype == cp.dtype("float32"):
+        cpp_make_blobs(
+            h_ptr[0],
+            <float*>x_p,
+            <int64_t*>y_p,
+            <int64_t>n_samples,
+            <int64_t>n_features,
+            <int64_t>n_centers,
+            row_c,
+            <const float*>ctr_p,
+            <const float*>0,
+            <float>cluster_std,
+            <bool>shuffle,
+            <float>center_box_min,
+            <float>center_box_max,
+            <uint64_t>random_state,
+        )
+    elif dtype == cp.dtype("float64"):
+        cpp_make_blobs(
+            h_ptr[0],
+            <double*>x_p,
+            <int64_t*>y_p,
+            <int64_t>n_samples,
+            <int64_t>n_features,
+            <int64_t>n_centers,
+            row_c,
+            <const double*>ctr_p,
+            <const double*>0,
+            <double>cluster_std,
+            <bool>shuffle,
+            <double>center_box_min,
+            <double>center_box_max,
+            <uint64_t>random_state,
+        )
+    else:
+        raise ValueError("RAFT make_blobs only supports float32 and float64.")
+
+    return X, y.astype(dtype, copy=False)

python/cuml/cuml/datasets/blobs.py

@@ -5,6 +5,7 @@
 
 import numbers
 from collections.abc import Iterable
+from random import getrandbits
 
 import cupy as cp
 import numpy as np
@@ -70,6 +71,115 @@ def _get_centers(rs, centers, center_box, n_samples, n_features, dtype):
     return centers, n_centers
 
 
+def _make_blobs_raft(
+    n_samples,
+    n_features,
+    centers,
+    cluster_std,
+    center_box,
+    shuffle,
+    random_state,
+    return_centers,
+    order,
+    dtype,
+):
+    n_samples = int(n_samples)
+    n_features = int(n_features)
+    dt = cp.dtype(dtype)
+
+    if n_samples <= 0 or n_features <= 0:
+        raise ValueError("`n_samples` and `n_features` must be positive.")
+
+    if cluster_std < 0:
+        raise ValueError("`cluster_std` must be non-negative.")
+
+    gen_ctr = centers is None or isinstance(centers, numbers.Integral)
+
+    if centers is None:
+        n_ctr, ctr = 3, None
+    elif isinstance(centers, numbers.Integral):
+        n_ctr, ctr = int(centers), None
+        if n_ctr <= 0:
+            raise ValueError("`centers` must be greater than 0.")
+    else:
+        ctr = cp.asarray(centers, dtype=dt, order=order)
+
+        if ctr.ndim != 2:
+            raise ValueError("`centers` must be a 2D array.")
+        if ctr.shape[1] != n_features:
+            raise ValueError(
+                "Expected `n_features` to be equal to"
+                " the length of axis 1 of centers array"
+            )
+
+        n_ctr = ctr.shape[0]
+        if n_ctr == 0:
+            raise ValueError("`centers` must contain at least one center.")
+
+    if gen_ctr:
+        try:
+            lo, hi = center_box
+        except (TypeError, ValueError):
+            raise ValueError(
+                "`center_box` must contain exactly two values."
+            ) from None
+
+        if lo > hi:
+            raise ValueError(
+                "`center_box` minimum must not exceed its maximum."
+            )
+    else:
+        lo = hi = 0.0
+
+    if return_centers and gen_ctr:
+        rs = _create_rs_generator(random_state=random_state)
+        ctr, n_ctr = _get_centers(
+            rs,
+            centers,
+            center_box,
+            n_samples,
+            n_features,
+            dt,
+        )
+
+    if ctr is not None:
+        ctr = cp.asarray(ctr, dtype=dt, order=order)
+
+        if order == "C" and not ctr.flags["C_CONTIGUOUS"]:
+            ctr = cp.ascontiguousarray(ctr)
+        elif order == "F" and not ctr.flags["F_CONTIGUOUS"]:
+            ctr = cp.asfortranarray(ctr)
+
+    if random_state is None:
+        seed = getrandbits(64)
+    else:
+        seed = int(random_state)
+
+        if not 0 <= seed <= (1 << 64) - 1:
+            raise ValueError("`random_state` must be between 0 and 2**64 - 1.")
+
+    from cuml.datasets._blobs import make_blobs as cpp_blobs
+
+    X, y = cpp_blobs(
+        n_samples=n_samples,
+        n_features=n_features,
+        n_centers=n_ctr,
+        centers=ctr,
+        cluster_std=float(cluster_std),
+        center_box_min=float(lo),
+        center_box_max=float(hi),
+        shuffle=bool(shuffle),
+        random_state=seed,
+        order=order,
+        dtype=dt,
+    )
+
+    if return_centers:
+        return X, y, ctr if gen_ctr else centers
+
+    return X, y
+
+
 @nvtx.annotate(message="datasets.make_blobs", domain="cuml_python")
 @cuml.internals.mlfunc(array_arg=None)
 def make_blobs(
@@ -151,6 +261,34 @@ def make_blobs(
     --------
     make_classification: a more intricate variant
     """
+    dt = cp.dtype(dtype)
+
+    use_cpp = (
+        isinstance(n_samples, numbers.Integral)
+        and isinstance(n_features, numbers.Integral)
+        and n_samples > 0
+        and n_features > 0
+        and isinstance(cluster_std, numbers.Real)
+        and isinstance(random_state, (type(None), int))
+        and shuffle is True
+        and order in ("C", "F")
+        and dt in (cp.dtype("float32"), cp.dtype("float64"))
+    )
+
+    if use_cpp:
+        return _make_blobs_raft(
+            n_samples=n_samples,
+            n_features=n_features,
+            centers=centers,
+            cluster_std=cluster_std,
+            center_box=center_box,
+            shuffle=shuffle,
+            random_state=random_state,
+            return_centers=return_centers,
+            order=order,
+            dtype=dt,
+        )
+
     generator = _create_rs_generator(random_state=random_state)
 
     centers, n_centers = _get_centers(