NVIDIA · switch527 · May 31, 2026
@@ -18,9 +18,7 @@ if(BUILD_CUML_BENCH)
     sg/linkage.cu
     sg/main.cpp
     sg/rf_classifier.cu
-    # FIXME: RF Regressor is having an issue where the tests now seem to take
-    # forever to finish, as opposed to the classifier counterparts!
-    # sg/rf_regressor.cu
+    sg/rf_regressor.cu
     sg/svc.cu
     sg/svr.cu
     sg/umap.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -52,12 +52,12 @@ class RFRegressor : public RegressionFixture<D> {
       auto* mPtr = &model.model;
       fit(*this->handle,
           mPtr,
-          this->data.X,
+          this->data.X.data(),
           this->params.nrows,
           this->params.ncols,
-          this->data.y,
+          this->data.y.data(),
           rfParams);
-      handle->sync_stream(this->stream);
+      this->handle->sync_stream(this->stream);
     });
   }
 
@@ -75,40 +75,38 @@ std::vector<RegParams> getInputs()
   struct std::vector<RegParams> out;
   RegParams p;
   p.data.rowMajor = false;
-  p.regression    = {.shuffle        = true,  // Better to shuffle when n_informative < ncols
-                     .effective_rank = -1,    // dataset generation will be faster
+  p.regression    = {.effective_rank = -1,  // dataset generation will be faster
                      .bias           = 4.5,
                      .tail_strength  = 0.5,  // unused when effective_rank = -1
                      .noise          = 1.0,
+                     .shuffle        = true,  // Better to shuffle when n_informative < ncols
                      .seed           = 12345ULL};
 
-  p.rf                          = set_rf_params(10,                 /*max_depth */
+  p.rf = set_rf_params(8,                  /*max_depth */
                        (1 << 20),          /* max_leaves */
                        0.3,                /* max_features */
                        32,                 /* max_n_bins */
                        3,                  /* min_samples_leaf */
                        3,                  /* min_samples_split */
                        0.0f,               /* min_impurity_decrease */
                        true,               /* bootstrap */
-                       500,                /* n_trees */
+                       100,                /* n_trees */
                        1.f,                /* max_samples */
                        1234ULL,            /* seed */
                        ML::CRITERION::MSE, /* split_criterion */
                        8,                  /* n_streams */
                        128                 /* max_batch_size */
   );
-  std::vector<DimInfo> dim_info = {{500000, 500, 400}};
+  // Scaled down from {500000, 500, 400} so the bench finishes in minutes;
+  // restore on a follow-up that addresses the FIXME.
+  std::vector<DimInfo> dim_info = {{100000, 100, 80}};
   for (auto& di : dim_info) {
-    // Let's run Bosch only for float type
-    if (!std::is_same<D, float>::value && di.ncols == 968) continue;
     p.data.nrows                  = di.nrows;
     p.data.ncols                  = di.ncols;
     p.regression.n_informative    = di.n_informative;
     p.rf.tree_params.max_features = 1.f;
-    for (auto max_depth : std::vector<int>({7, 11, 15})) {
-      p.rf.tree_params.max_depth = max_depth;
-      out.push_back(p);
-    }
+    p.rf.tree_params.max_depth    = 8;
+    out.push_back(p);
   }
   return out;
 }

@@ -151,7 +151,8 @@ void fit(const raft::handle_t& user_handle,
          int n_unique_labels,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
-         bool* bootstrap_masks               = nullptr);
+         bool* bootstrap_masks               = nullptr,
+         float* sample_weight                = nullptr);
 void fit(const raft::handle_t& user_handle,
          RandomForestClassifierD* forest,
          double* input,
@@ -161,7 +162,8 @@ void fit(const raft::handle_t& user_handle,
          int n_unique_labels,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
-         bool* bootstrap_masks               = nullptr);
+         bool* bootstrap_masks               = nullptr,
+         double* sample_weight               = nullptr);
 
 template <typename T, typename L>
 void fit_treelite(const raft::handle_t& user_handle,
@@ -174,7 +176,8 @@ void fit_treelite(const raft::handle_t& user_handle,
                   RF_params rf_params,
                   bool* bootstrap_masks,
                   T* feature_importances,
-                  rapids_logger::level_enum verbosity);
+                  rapids_logger::level_enum verbosity,
+                  T* sample_weight = nullptr);
 
 void predict(const raft::handle_t& user_handle,
              const RandomForestClassifierF* forest,
@@ -232,7 +235,8 @@ void fit(const raft::handle_t& user_handle,
          float* labels,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
-         bool* bootstrap_masks               = nullptr);
+         bool* bootstrap_masks               = nullptr,
+         float* sample_weight                = nullptr);
 void fit(const raft::handle_t& user_handle,
          RandomForestRegressorD* forest,
          double* input,
@@ -241,7 +245,8 @@ void fit(const raft::handle_t& user_handle,
          double* labels,
          RF_params rf_params,
          rapids_logger::level_enum verbosity = rapids_logger::level_enum::info,
-         bool* bootstrap_masks               = nullptr);
+         bool* bootstrap_masks               = nullptr,
+         double* sample_weight               = nullptr);
 
 template <typename T, typename L>
 void fit_treelite(const raft::handle_t& user_handle,
@@ -253,7 +258,8 @@ void fit_treelite(const raft::handle_t& user_handle,
                   RF_params rf_params,
                   bool* bootstrap_masks,
                   T* feature_importances,
-                  rapids_logger::level_enum verbosity);
+                  rapids_logger::level_enum verbosity,
+                  T* sample_weight = nullptr);
 
 void predict(const raft::handle_t& user_handle,
              const RandomForestRegressorF* forest,

@@ -9,17 +9,17 @@ namespace ML {
 namespace DT {
 
 struct CountBin {
-  // double covers both the unweighted count path and the future weighted-count
-  // path with one bin type; 32-bit int would overflow on large weighted counts.
+  // double covers both the unweighted count path (weight=1.0) and the
+  // per-sample-weighted path; 32-bit int would overflow on large weighted counts.
   double x;
   CountBin(CountBin const&) = default;
   HDI CountBin(double x_) : x(x_) {}
   HDI CountBin() : x(0.0) {}
 
-  DI static void IncrementHistogram(CountBin* hist, int n_bins, int b, int label)
+  DI static void IncrementHistogram(CountBin* hist, int n_bins, int b, int label, double weight)
   {
     auto offset = label * n_bins + b;
-    CountBin::AtomicAdd(hist + offset, {1.0});
+    CountBin::AtomicAdd(hist + offset, {weight});
   }
   DI static void AtomicAdd(CountBin* address, CountBin val) { atomicAdd(&address->x, val.x); }
   HDI CountBin& operator+=(const CountBin& b)
@@ -35,16 +35,19 @@ struct CountBin {
 };
 
 struct AggregateBin {
+  // `label_sum` carries `label * weight`; `count` stays unweighted so
+  // `min_samples_leaf` checks operate on the integer sample count.
   double label_sum;
   int count;
 
   AggregateBin(AggregateBin const&) = default;
   HDI AggregateBin() : label_sum(0.0), count(0) {}
   HDI AggregateBin(double label_sum, int count) : label_sum(label_sum), count(count) {}
 
-  DI static void IncrementHistogram(AggregateBin* hist, int n_bins, int b, double label)
+  DI static void IncrementHistogram(
+    AggregateBin* hist, int n_bins, int b, double label, double weight)
   {
-    AggregateBin::AtomicAdd(hist + b, {label, 1});
+    AggregateBin::AtomicAdd(hist + b, {label * weight, 1});
   }
   DI static void AtomicAdd(AggregateBin* address, AggregateBin val)
   {
@@ -63,5 +66,6 @@ struct AggregateBin {
     return b;
   }
 };
+static_assert(sizeof(AggregateBin) == 16, "AggregateBin layout drift");
 }  // namespace DT
 }  // namespace ML
@@ -162,6 +162,17 @@ struct Builder {
   IdxT* n_nodes;
   /** buffer of segmented histograms*/
   BinT* histograms;
+  /** Classifier companion: per-bin unweighted sample counts for
+   *  `min_samples_leaf` + `Split::nLeft` under fractional weights. */
+  int* unweighted_histograms = nullptr;
+  /** Regressor companion: per-bin sum-of-weights totals for the weighted
+   *  gain denominator and leaf mean. */
+  double* weighted_count_histograms = nullptr;
+  // Compile-time gates on the builder's BinT type; mutually exclusive.
+  static constexpr bool kIsClassifier = std::is_same<typename ObjectiveT::BinT, CountBin>::value;
+  static constexpr bool kIsRegressor = std::is_same<typename ObjectiveT::BinT, AggregateBin>::value;
+  static_assert(kIsClassifier || kIsRegressor,
+                "Builder<ObjectiveT>::BinT must be CountBin or AggregateBin");
   /** threadblock arrival count */
   int* done_count;
   /** mutex array used for atomically updating best split */
@@ -199,7 +210,8 @@ struct Builder {
           IdxT n_cols,
           rmm::device_uvector<IdxT>* row_ids,
           IdxT n_classes,
-          const QuantilesT& q)
+          const QuantilesT& q,
+          const DataT* sample_weight = nullptr)
     : handle(handle),
       builder_stream(s),
       treeid(treeid),
@@ -212,7 +224,8 @@ struct Builder {
               int(row_ids->size()),
               max(1, IdxT(params.max_features * n_cols)),
               row_ids->data(),
-              n_classes},
+              n_classes,
+              sample_weight},
       quantiles(q),
       d_buff(0, builder_stream)
   {
@@ -266,8 +279,15 @@ struct Builder {
     size_t max_len_histograms =
       max_batch * params.max_n_bins * n_blks_for_cols * dataset.num_outputs;
 
-    d_wsize += calculateAlignedBytes(sizeof(IdxT));                               // n_nodes
-    d_wsize += calculateAlignedBytes(sizeof(BinT) * max_len_histograms);          // histograms
+    d_wsize += calculateAlignedBytes(sizeof(IdxT));                       // n_nodes
+    d_wsize += calculateAlignedBytes(sizeof(BinT) * max_len_histograms);  // histograms
+    size_t max_len_companion = max_batch * params.max_n_bins * n_blks_for_cols;
+    if constexpr (kIsClassifier) {
+      d_wsize += calculateAlignedBytes(sizeof(int) * max_len_companion);  // unweighted_histograms
+    } else if constexpr (kIsRegressor) {
+      d_wsize +=
+        calculateAlignedBytes(sizeof(double) * max_len_companion);  // weighted_count_histograms
+    }
     d_wsize += calculateAlignedBytes(sizeof(int) * max_batch * n_blks_for_cols);  // done_count
     d_wsize += calculateAlignedBytes(sizeof(int) * max_batch);                    // mutex
     d_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch);                 // splits
@@ -304,6 +324,14 @@ struct Builder {
     d_wspace += calculateAlignedBytes(sizeof(IdxT));
     histograms = reinterpret_cast<BinT*>(d_wspace);
     d_wspace += calculateAlignedBytes(sizeof(BinT) * max_len_histograms);
+    size_t max_len_companion = max_batch * (params.max_n_bins) * n_blks_for_cols;
+    if constexpr (kIsClassifier) {
+      unweighted_histograms = reinterpret_cast<int*>(d_wspace);
+      d_wspace += calculateAlignedBytes(sizeof(int) * max_len_companion);
+    } else if constexpr (kIsRegressor) {
+      weighted_count_histograms = reinterpret_cast<double*>(d_wspace);
+      d_wspace += calculateAlignedBytes(sizeof(double) * max_len_companion);
+    }
     done_count = reinterpret_cast<int*>(d_wspace);
     d_wspace += calculateAlignedBytes(sizeof(int) * max_batch * n_col_blks);
     mutex = reinterpret_cast<int*>(d_wspace);
@@ -487,12 +515,20 @@ struct Builder {
   auto computeSplitSmemSize()
   {
     size_t smem_size_1 =
-      params.max_n_bins * dataset.num_outputs * sizeof(BinT) +  // shared_histogram size
-      params.max_n_bins * sizeof(DataT) +                       // shared_quantiles size
-      sizeof(int);                                              // shared_done size
-    // Extra room for alignment (see alignPointer in
-    // computeSplitKernel)
-    smem_size_1 += sizeof(DataT) + 3 * sizeof(int);
+      params.max_n_bins * dataset.num_outputs * sizeof(BinT) +  // shared_histogram
+      params.max_n_bins * sizeof(DataT) +                       // shared_quantiles
+      sizeof(int);                                              // shared_done
+    int n_align_slots = 3;  // shared_histogram, shared_quantiles, shared_done
+    if constexpr (kIsClassifier) {
+      smem_size_1 += params.max_n_bins * sizeof(int);  // shared_unweighted
+      n_align_slots = 4;
+    } else if constexpr (kIsRegressor) {
+      smem_size_1 += params.max_n_bins * sizeof(double);  // shared_weighted_count
+      n_align_slots = 4;
+    }
+    // Worst-case alignPointer slack per slot is sizeof(largest-following-type)-1.
+    // The regressor companion is double-aligned, so use 8 bytes per slot.
+    smem_size_1 += n_align_slots * sizeof(double);
     // Calculate the shared memory needed for evalBestSplit
     size_t smem_size_2 = raft::ceildiv(TPB_DEFAULT, raft::WarpSize) * sizeof(SplitT);
     // Pick the max of two
@@ -517,12 +553,22 @@ struct Builder {
     // required total length (in bins) of the global segmented histograms over all
     // classes, features and (large)nodes.
     int len_histograms = n_bins * n_classes * n_blocks_dimy * n_large_nodes;
+    int len_companion  = n_bins * n_blocks_dimy * n_large_nodes;
     RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream));
+    if constexpr (kIsClassifier) {
+      RAFT_CUDA_TRY(
+        cudaMemsetAsync(unweighted_histograms, 0, sizeof(int) * len_companion, builder_stream));
+    } else if constexpr (kIsRegressor) {
+      RAFT_CUDA_TRY(cudaMemsetAsync(
+        weighted_count_histograms, 0, sizeof(double) * len_companion, builder_stream));
+    }
     // create the objective function object
     ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf);
     // call the computeSplitKernel
     raft::common::nvtx::range kernel_scope("computeSplitKernel @builder.cuh [batched-levelalgo]");
     launchComputeSplitKernel<DataT, LabelT, IdxT, TPB_DEFAULT>(histograms,
+                                                               unweighted_histograms,
+                                                               weighted_count_histograms,
                                                                params.max_n_bins,
                                                                params.min_samples_split,
                                                                params.max_leaves,

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -26,6 +26,8 @@ struct Dataset {
   IdxT* row_ids;
   /** Number of classes or regression outputs*/
   IdxT num_outputs;
+  /** per-row sample weights (length M); nullptr when no per-sample weighting */
+  const DataT* sample_weight = nullptr;
 };
 
 }  // namespace DT

@@ -365,13 +365,17 @@ CUML_KERNEL void adaptive_sample_kernel(int* colids,
   }
 }
 
+// Exactly one of the two companion pointers is non-null per builder; reads
+// are gated on `kIsClassifier` / `kIsRegressor`.
 template <typename DataT,
           typename LabelT,
           typename IdxT,
           int TPB,
           typename ObjectiveT,
           typename BinT>
 void launchComputeSplitKernel(BinT* histograms,
+                              int* unweighted_histograms,
+                              double* weighted_count_histograms,
                               IdxT n_bins,
                               IdxT min_samples_split,
                               IdxT max_leaves,