Refactor kernel_density to use less memory

cpp/CMakeLists.txt

@@ -427,7 +427,7 @@ if(BUILD_CUML_CPP_LIBRARY)
   endif()
 
   if(all_algo OR knn_algo)
-    target_sources(cuml_objs PRIVATE src/knn/knn.cu src/knn/knn_sparse.cu)
+    target_sources(cuml_objs PRIVATE src/knn/knn.cu src/knn/knn_sparse.cu src/kde/kde.cu)
   endif()
 
   if(all_algo OR hierarchicalclustering_algo)

cpp/include/cuml/neighbors/kde.hpp

@@ -0,0 +1,90 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <cuml/common/distance_type.hpp>
+
+#include <raft/core/resources.hpp>
+
+#include <cstdint>
+
+namespace ML::KDE {
+
+enum class DensityKernelType : int {
+  Gaussian     = 0,
+  Tophat       = 1,
+  Epanechnikov = 2,
+  Exponential  = 3,
+  Linear       = 4,
+  Cosine       = 5
+};
+
+/**
+ * @brief Compute normalized log-density scores for query samples.
+ *
+ * The query and training arrays must be dense row-major (C-contiguous)
+ * device arrays with shapes `(n_query, n_features)` and
+ * `(n_train, n_features)`, respectively.
+ *
+ * @tparam T floating point type, either float or double
+ * @param[in] handle raft resources used to launch work
+ * @param[in] query device pointer to query samples in row-major order
+ * @param[in] train device pointer to training samples in row-major order
+ * @param[in] weights optional device pointer to sample weights of length
+ * `n_train`, or nullptr for uniform weights
+ * @param[out] output device pointer to log-density scores of length `n_query`
+ * @param[in] n_query number of query samples
+ * @param[in] n_train number of training samples
+ * @param[in] n_features number of features per sample
+ * @param[in] bandwidth positive KDE bandwidth
+ * @param[in] sum_weights sum of `weights`, or `n_train` when weights is null
+ * @param[in] kernel density kernel to evaluate
+ * @param[in] metric distance metric used between query and training samples
+ * @param[in] metric_arg metric-specific argument, such as p for Minkowski
+ */
+template <typename T>
+void score_samples(raft::resources const& handle,
+                   const T* query,
+                   const T* train,
+                   const T* weights,
+                   T* output,
+                   std::int64_t n_query,
+                   std::int64_t n_train,
+                   std::int64_t n_features,
+                   T bandwidth,
+                   T sum_weights,
+                   DensityKernelType kernel,
+                   ML::distance::DistanceType metric,
+                   T metric_arg);
+
+extern template void score_samples<float>(raft::resources const&,
+                                          const float*,
+                                          const float*,
+                                          const float*,
+                                          float*,
+                                          std::int64_t,
+                                          std::int64_t,
+                                          std::int64_t,
+                                          float,
+                                          float,
+                                          DensityKernelType,
+                                          ML::distance::DistanceType,
+                                          float);
+
+extern template void score_samples<double>(raft::resources const&,
+                                           const double*,
+                                           const double*,
+                                           const double*,
+                                           double*,
+                                           std::int64_t,
+                                           std::int64_t,
+                                           std::int64_t,
+                                           double,
+                                           double,
+                                           DensityKernelType,
+                                           ML::distance::DistanceType,
+                                           double);
+
+}  // namespace ML::KDE

cpp/src/kde/kde.cu

@@ -0,0 +1,83 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cuml/neighbors/kde.hpp>
+
+#include <raft/core/device_mdspan.hpp>
+
+#include <cuvs/distance/kde.hpp>
+
+#include <optional>
+
+namespace ML::KDE {
+
+template <typename T>
+void score_samples(raft::resources const& handle,
+                   const T* query,
+                   const T* train,
+                   const T* weights,
+                   T* output,
+                   std::int64_t n_query,
+                   std::int64_t n_train,
+                   std::int64_t n_features,
+                   T bandwidth,
+                   T sum_weights,
+                   DensityKernelType kernel,
+                   ML::distance::DistanceType metric,
+                   T metric_arg)
+{
+  auto query_view =
+    raft::make_device_matrix_view<const T, std::int64_t>(query, n_query, n_features);
+  auto train_view =
+    raft::make_device_matrix_view<const T, std::int64_t>(train, n_train, n_features);
+  auto output_view = raft::make_device_vector_view<T, std::int64_t>(output, n_query);
+  auto weights_view =
+    weights
+      ? std::make_optional(raft::make_device_vector_view<const T, std::int64_t>(weights, n_train))
+      : std::nullopt;
+  auto cuvs_kernel = static_cast<cuvs::distance::DensityKernelType>(kernel);
+  auto cuvs_metric = static_cast<cuvs::distance::DistanceType>(metric);
+
+  cuvs::distance::kde(handle,
+                      query_view,
+                      train_view,
+                      weights_view,
+                      output_view,
+                      bandwidth,
+                      sum_weights,
+                      cuvs_kernel,
+                      cuvs_metric,
+                      metric_arg);
+}
+
+template void score_samples<float>(raft::resources const&,
+                                   const float*,
+                                   const float*,
+                                   const float*,
+                                   float*,
+                                   std::int64_t,
+                                   std::int64_t,
+                                   std::int64_t,
+                                   float,
+                                   float,
+                                   DensityKernelType,
+                                   ML::distance::DistanceType,
+                                   float);
+
+template void score_samples<double>(raft::resources const&,
+                                    const double*,
+                                    const double*,
+                                    const double*,
+                                    double*,
+                                    std::int64_t,
+                                    std::int64_t,
+                                    std::int64_t,
+                                    double,
+                                    double,
+                                    DensityKernelType,
+                                    ML::distance::DistanceType,
+                                    double);
+
+}  // namespace ML::KDE

python/cuml/cuml/neighbors/CMakeLists.txt

@@ -1,11 +1,12 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
 
 set(cython_sources "")
+add_module_gpu_default("kernel_density.pyx" ${knn_algo} ${neighbors_algo})
 add_module_gpu_default("kneighbors_classifier.pyx" ${kneighbors_classifier_algo} ${neighbors_algo})
 add_module_gpu_default("kneighbors_regressor.pyx" ${kneighbors_regressor_algo} ${neighbors_algo})
 add_module_gpu_default("nearest_neighbors.pyx" ${nearest_neighbors_algo} ${neighbors_algo})