From 7569118c4e3dfb1ee1ccc4456eca0bf9ce5a2388 Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Mon, 6 Jul 2026 10:23:47 +0200 Subject: [PATCH 1/4] Handle RF split histogram shared-memory pressure --- .../batched-levelalgo/builder.cuh | 103 ++++++++++++------ .../kernels/builder_kernels.cuh | 3 +- .../kernels/builder_kernels_impl.cuh | 94 ++++++++++------ .../kernels/classification-double.cu | 1 + .../kernels/classification-float.cu | 1 + .../kernels/regression-double.cu | 1 + .../kernels/regression-float.cu | 1 + .../kernels/weighted-classification-double.cu | 1 + .../kernels/weighted-classification-float.cu | 1 + .../kernels/weighted-regression-double.cu | 1 + .../kernels/weighted-regression-float.cu | 1 + cpp/tests/sg/rf_test.cu | 38 +++++++ 12 files changed, 173 insertions(+), 73 deletions(-) diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index 25a151fb92..628b4f4b8f 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -149,6 +149,11 @@ struct Builder { /** default threads per block for most kernels in here */ static constexpr int TPB_DEFAULT = 128; + // Tunable performance heuristic for the shared-memory histogram path. Large per-block + // histograms, usually from large n_classes, can reduce occupancy enough that global memory is + // faster even when the histogram fits in shared memory. 16 KiB keeps small/default histograms in + // shared memory while avoiding the large-class shared-memory slowdown measured locally. + static constexpr size_t tunable_split_histogram_dynamic_smem_limit_bytes = 16 * 1024; /** handle to get device properties */ const raft::handle_t& handle; /** stream to launch kernels */ @@ -491,7 +496,7 @@ struct Builder { sampleFeatures(work_items, sampling_seed, sample_offset); for (IdxT c = 0; c < dataset.n_sampled_cols; c += n_blks_for_cols) { - computeSplit(c, n_blocks_dimx, n_large_nodes); + computeSplit(c, n_blocks_dimx, n_large_nodes, work_items.size()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } raft::update_host(h_splits, splits, work_items.size(), builder_stream); @@ -515,7 +520,7 @@ struct Builder { RAFT_CUDA_TRY(cudaPeekAtLastError()); } - auto computeSplitSmemSize() + size_t computeSplitHistogramSmemSize() const { auto shared_histogram_size = ML::checked_mul(params.max_n_bins, dataset.num_outputs, sizeof(BinT)); @@ -527,60 +532,86 @@ struct Builder { // computeSplitKernel) auto alignment_smem_size = ML::checked_add(sizeof(DataT), ML::checked_mul(3, sizeof(int))); - dynamic_smem_size = ML::checked_add(dynamic_smem_size, alignment_smem_size); + return ML::checked_add(dynamic_smem_size, alignment_smem_size); + } + + size_t computeSplitGlobalHistogramSmemSize() const + { + // shared_done only, plus conservative alignment room for alignPointer. + return ML::checked_add(sizeof(int), sizeof(int)); + } + size_t computeSplitStaticSmemSize() const + { // computeSplitKernel also reserves static shared memory for CUB's scan temp // storage and the per-warp split reduction scratch. auto cdf_scan_smem_size = sizeof(typename cub::BlockScan::TempStorage); auto split_scratch_smem_size = ML::checked_mul(raft::ceildiv(TPB_DEFAULT, raft::WarpSize), sizeof(SplitT)); - auto total_smem_size = - ML::checked_add(dynamic_smem_size, cdf_scan_smem_size, split_scratch_smem_size); - auto available_smem = handle.get_device_properties().sharedMemPerBlock; - ASSERT(available_smem >= total_smem_size, - "Not enough shared memory. Consider reducing max_n_bins."); - return dynamic_smem_size; + return ML::checked_add(cdf_scan_smem_size, split_scratch_smem_size); + } + + size_t computeSplitSmemSize() const + { + return ML::checked_add(computeSplitHistogramSmemSize(), + computeSplitStaticSmemSize()); + } + + bool shouldUseGlobalMemoryHistogram(size_t shared_histogram_dynamic_smem_size, + size_t shared_path_total_smem_size) const + { + auto available_smem = size_t(handle.get_device_properties().sharedMemPerBlock); + auto global_smem = ML::checked_add(computeSplitGlobalHistogramSmemSize(), + computeSplitStaticSmemSize()); + ASSERT(available_smem >= global_smem, "Not enough shared memory for RF split bookkeeping."); + return shared_path_total_smem_size > available_smem || + shared_histogram_dynamic_smem_size > tunable_split_histogram_dynamic_smem_limit_bytes; } - void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes) + void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes, size_t n_work_items) { // if no instances to split, return if (n_blocks_dimx == 0) return; raft::common::nvtx::range fun_scope("Builder::computeSplit @builder.cuh [batched-levelalgo]"); - auto n_bins = params.max_n_bins; - auto n_classes = dataset.num_outputs; + auto n_bins = params.max_n_bins; + auto n_classes = dataset.num_outputs; + auto shared_histogram_dynamic_smem_size = computeSplitHistogramSmemSize(); + auto shared_path_total_smem_size = computeSplitSmemSize(); + auto use_global_memory_histogram = shouldUseGlobalMemoryHistogram( + shared_histogram_dynamic_smem_size, shared_path_total_smem_size); + auto smem_size = use_global_memory_histogram ? computeSplitGlobalHistogramSmemSize() + : shared_histogram_dynamic_smem_size; // if columns left to be processed lesser than `n_blks_for_cols`, shrink the blocks along dimy auto n_blocks_dimy = std::min(n_blks_for_cols, dataset.n_sampled_cols - col); - // compute required dynamic shared memory - auto smem_size = computeSplitSmemSize(); dim3 grid(n_blocks_dimx, n_blocks_dimy, 1); - // 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; + auto histogram_node_count = use_global_memory_histogram ? n_work_items : n_large_nodes; + size_t len_histograms = size_t(n_bins) * n_classes * n_blocks_dimy * histogram_node_count; RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream)); // create the objective function object ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf, params.split_criterion); // call the computeSplitKernel raft::common::nvtx::range kernel_scope("computeSplitKernel @builder.cuh [batched-levelalgo]"); - launchComputeSplitKernel(histograms, - params.max_n_bins, - params.min_samples_split, - params.max_leaves, - dataset, - quantiles, - d_work_items, - col, - column_samples, - done_count, - mutex, - splits, - objective, - treeid, - workload_info, - seed, - grid, - smem_size, - builder_stream); + launchComputeSplitKernel( + histograms, + params.max_n_bins, + params.min_samples_split, + params.max_leaves, + dataset, + quantiles, + d_work_items, + col, + column_samples, + done_count, + mutex, + splits, + objective, + treeid, + workload_info, + seed, + use_global_memory_histogram, + grid, + smem_size, + builder_stream); } // Set the leaf value predictions in batch diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh index c0181fbaec..b13ba2d9d0 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh @@ -126,7 +126,7 @@ void launchLeafKernel(ObjectiveT objective, // Values outside the quantile range are clamped to the edge bins: values below the // first quantile return 0, and values above the last quantile return len - 1. template -HDI IdxT lower_bound(DataT* array, IdxT len, DataT element) +HDI IdxT lower_bound(DataT const* array, IdxT len, DataT element) { IdxT start = 0; IdxT end = len - 1; @@ -159,6 +159,7 @@ void launchComputeSplitKernel(typename ObjectiveT::BinT* histograms, IdxT treeid, const WorkloadInfo* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh index fb7c70e9c1..faf0fe241f 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh @@ -223,14 +223,14 @@ void launchLeafKernel(ObjectiveT objective, } /** - * @brief For every threadblock, converts the smem pdf-histogram to + * @brief For every threadblock, converts a pdf-histogram to a * cdf-histogram inplace using inclusive block-sum-scan and returns * the total_sum * @return The total sum aggregated over the sumscan, * as well as the modified cdf-histogram pointer */ template -DI BinT pdf_to_cdf(BinT* shared_histogram, IdxT n_bins) +DI BinT pdf_to_cdf(BinT* histogram, IdxT n_bins) { // Blockscan instance preparation typedef cub::BlockScan BlockScan; @@ -242,10 +242,10 @@ DI BinT pdf_to_cdf(BinT* shared_histogram, IdxT n_bins) for (IdxT tix = threadIdx.x; tix < raft::ceildiv(n_bins, TPB) * TPB; tix += blockDim.x) { BinT result; BinT block_aggregate; - BinT element = tix < n_bins ? shared_histogram[tix] : BinT(); + BinT element = tix < n_bins ? histogram[tix] : BinT(); BlockScan(temp_storage).InclusiveSum(element, result, block_aggregate); __syncthreads(); - if (tix < n_bins) { shared_histogram[tix] = result + total_aggregate; } + if (tix < n_bins) { histogram[tix] = result + total_aggregate; } total_aggregate += block_aggregate; } // return the total sum @@ -268,7 +268,8 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms, ObjectiveT objective, IdxT treeid, const WorkloadInfo* workload_info, - uint64_t seed) + uint64_t seed, + bool use_global_memory_histogram) { using BinT = typename ObjectiveT::BinT; // dynamic shared memory @@ -296,24 +297,35 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms, // getting the n_bins for that feature int n_bins = quantiles.n_bins_array[col]; + auto n_classes = objective.NumClasses(); auto end = range_start + range_len; - auto shared_histogram_len = n_bins * objective.NumClasses(); - auto* shared_histogram = alignPointer(smem); - auto* shared_quantiles = alignPointer(shared_histogram + shared_histogram_len); - auto* shared_done = alignPointer(shared_quantiles + n_bins); + auto histogram_len = n_bins * n_classes; + auto* histogram = static_cast(nullptr); + auto* shared_done = static_cast(nullptr); + auto* quantiles_for_split = quantiles.quantiles_array + std::size_t(max_n_bins) * col; IdxT stride = blockDim.x * num_blocks; IdxT tid = threadIdx.x + offset_blockid * blockDim.x; - // populating shared memory with initial values - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) - shared_histogram[i] = BinT(); - for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) - shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b]; + if (use_global_memory_histogram) { + auto histograms_offset = (std::size_t(nid) * gridDim.y + blockIdx.y) * max_n_bins * n_classes; + histogram = histograms + histograms_offset; + shared_done = alignPointer(smem); + } else { + histogram = alignPointer(smem); + auto* shared_quantiles = alignPointer(histogram + histogram_len); + shared_done = alignPointer(shared_quantiles + n_bins); + quantiles_for_split = shared_quantiles; + for (IdxT i = threadIdx.x; i < histogram_len; i += blockDim.x) { + histogram[i] = BinT(); + } + for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) { + shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b]; + } + } - // synchronizing above changes across block __syncthreads(); - // compute pdf shared histogram for all bins for all classes in shared mem + // compute pdf histogram for all bins for all classes // Must be 64 bit - can easily grow larger than a 32 bit int std::size_t col_offset = std::size_t(col) * dataset.M; @@ -323,20 +335,28 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms, auto data = dataset.data[row + col_offset]; auto label = dataset.labels[row]; - // `start` is lowest index such that data <= shared_quantiles[start] - IdxT start = lower_bound(shared_quantiles, n_bins, data); - // ++shared_histogram[start] - objective.IncrementHistogram(shared_histogram, n_bins, start, label, dataset, row); + // `start` is lowest index such that data <= quantiles_for_split[start] + IdxT start = lower_bound(quantiles_for_split, n_bins, data); + // ++histogram[start] + objective.IncrementHistogram(histogram, n_bins, start, label, dataset, row); } - // synchronizing above changes across block __syncthreads(); - if (num_blocks > 1) { + if (use_global_memory_histogram) { + __threadfence(); // for commit guarantee before the last block scores the split + __syncthreads(); + + bool last = MLCommon::signalDone( + done_count + nid * gridDim.y + blockIdx.y, num_blocks, offset_blockid == 0, shared_done); + if (!last) return; + } else if (num_blocks > 1) { + // Shared-memory histogram path: each block built a partial histogram, so unify those + // partial histograms in global memory before scoring the split. // update the corresponding global location auto histograms_offset = - ((large_nid * gridDim.y) + blockIdx.y) * max_n_bins * objective.NumClasses(); - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { - BinT::AtomicAdd(histograms + histograms_offset + i, shared_histogram[i]); + (std::size_t(large_nid) * gridDim.y + blockIdx.y) * max_n_bins * n_classes; + for (IdxT i = threadIdx.x; i < histogram_len; i += blockDim.x) { + BinT::AtomicAdd(histograms + histograms_offset + i, histogram[i]); } __threadfence(); // for commit guarantee @@ -349,34 +369,34 @@ static __global__ void computeSplitKernel(typename ObjectiveT::BinT* histograms, if (!last) return; // store the complete global histogram in shared memory of last block - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) - shared_histogram[i] = histograms[histograms_offset + i]; + for (IdxT i = threadIdx.x; i < histogram_len; i += blockDim.x) { + histogram[i] = histograms[histograms_offset + i]; + } __syncthreads(); } - // PDF to CDF inplace in `shared_histogram` - for (IdxT c = 0; c < objective.NumClasses(); ++c) { + // PDF to CDF inplace in `histogram` + for (IdxT c = 0; c < n_classes; ++c) { // left to right scan operation for scanning // "lesser-than-or-equal" counts - BinT total_sum = pdf_to_cdf(shared_histogram + n_bins * c, n_bins); - // now, `shared_histogram[n_bins * c + i]` will have count of datapoints of class `c` - // that are less than or equal to `shared_quantiles[i]`. + BinT total_sum = pdf_to_cdf(histogram + n_bins * c, n_bins); + // now, `histogram[n_bins * c + i]` will have count of datapoints of class `c` + // that are less than or equal to `quantiles_for_split[i]`. } __syncthreads(); // calculate the best candidate bins (one for each thread in the block) in current feature and // corresponding information gain for splitting - Split sp = - objective.Gain(shared_histogram, shared_quantiles, col, range_len, n_bins); + Split sp = objective.Gain(histogram, quantiles_for_split, col, range_len, n_bins); __syncthreads(); // calculate best bins among candidate bins per feature using warp reduce // then atomically update across features to get best split per node // (in split[nid]) - sp.evalBestSplit(split_scratch, splits + nid, mutex + nid, shared_quantiles, n_bins); + sp.evalBestSplit(split_scratch, splits + nid, mutex + nid, quantiles_for_split, n_bins); } template @@ -396,6 +416,7 @@ void launchComputeSplitKernel(typename ObjectiveT::BinT* histograms, IdxT treeid, const WorkloadInfo* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream) @@ -416,7 +437,8 @@ void launchComputeSplitKernel(typename ObjectiveT::BinT* histograms, objective, treeid, workload_info, - seed); + seed, + use_global_memory_histogram); } } // namespace DT diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu index e050054f26..919429375c 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu index 6278efd092..32560e72cf 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu index e42c1bd0cd..b9b72a6495 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu index 42780bd796..68d888fd97 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu index 69e8d3af1c..2d9590d097 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu index 0cb06c3e01..f74ff643ac 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu index c735997614..442cc9c024 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu index 11015f4da0..115d00af22 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu @@ -46,6 +46,7 @@ template void launchComputeSplitKernel* workload_info, uint64_t seed, + bool use_global_memory_histogram, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/tests/sg/rf_test.cu b/cpp/tests/sg/rf_test.cu index a2a08572d4..bac61ace5e 100644 --- a/cpp/tests/sg/rf_test.cu +++ b/cpp/tests/sg/rf_test.cu @@ -842,6 +842,44 @@ TEST(RfTests, IntegerOverflow) handle.sync_stream_pool(); } +TEST(RfTests, HighClassCountSplitHistogramFallsBackToGlobalMemory) +{ + constexpr std::size_t n_rows = 640; + constexpr std::size_t n_cols = 4; + constexpr int n_classes = 80; + constexpr int max_n_bins = 256; + + auto stream_pool = std::make_shared(1); + raft::handle_t handle(rmm::cuda_stream_per_thread, stream_pool); + thrust::device_vector X(n_rows * n_cols); + thrust::device_vector y(n_rows); + + Datasets::make_blobs(handle, + X.data().get(), + y.data().get(), + n_rows, + n_cols, + n_classes, + false, + nullptr, + nullptr, + 5.0, + false, + -10.0f, + 10.0f, + 1234); + + auto forest = std::make_shared>(); + auto rf_params = + set_rf_params(2, -1, 1.0f, max_n_bins, 1, 2, 0.0f, false, 1, 1.0f, 1234, CRITERION::GINI, 1, 4); + auto forest_ptr = forest.get(); + + ASSERT_NO_THROW( + fit(handle, forest_ptr, X.data().get(), n_rows, n_cols, y.data().get(), n_classes, rf_params)); + EXPECT_EQ(forest->trees.size(), 1); + EXPECT_GE(forest->trees[0]->leaf_counter, 1); +} + TEST(RfTests, InvalidSampleWeightThrows) { constexpr std::size_t n_rows = 16; From 8f2b98a1ef9a1bbfa6c667f7d096f69f3309bc9a Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Mon, 6 Jul 2026 10:37:11 +0200 Subject: [PATCH 2/4] Strengthen RF global histogram fallback test --- cpp/tests/sg/rf_test.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tests/sg/rf_test.cu b/cpp/tests/sg/rf_test.cu index bac61ace5e..999cafdbec 100644 --- a/cpp/tests/sg/rf_test.cu +++ b/cpp/tests/sg/rf_test.cu @@ -876,8 +876,8 @@ TEST(RfTests, HighClassCountSplitHistogramFallsBackToGlobalMemory) ASSERT_NO_THROW( fit(handle, forest_ptr, X.data().get(), n_rows, n_cols, y.data().get(), n_classes, rf_params)); - EXPECT_EQ(forest->trees.size(), 1); - EXPECT_GE(forest->trees[0]->leaf_counter, 1); + ASSERT_EQ(forest->trees.size(), 1); + EXPECT_GT(forest->trees.front()->depth_counter, 0); } TEST(RfTests, InvalidSampleWeightThrows) From 3bcf2b38ddcd79906cfea875a6912c7f47a0d3cc Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Mon, 6 Jul 2026 12:12:31 +0200 Subject: [PATCH 3/4] Hoist RF split shared-memory config --- .../batched-levelalgo/builder.cuh | 87 ++++++++++--------- 1 file changed, 44 insertions(+), 43 deletions(-) diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index 628b4f4b8f..567eae2768 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -492,11 +492,12 @@ struct Builder { RAFT_CUDA_TRY(cudaMemsetAsync(mutex, 0, sizeof(int) * params.max_batch_size, builder_stream)); raft::update_device(d_work_items, work_items.data(), work_items.size(), builder_stream); auto [n_blocks_dimx, n_large_nodes] = this->updateWorkloadInfo(work_items); + auto split_smem_config = computeSplitSharedMemoryConfig(); sampleFeatures(work_items, sampling_seed, sample_offset); for (IdxT c = 0; c < dataset.n_sampled_cols; c += n_blks_for_cols) { - computeSplit(c, n_blocks_dimx, n_large_nodes, work_items.size()); + computeSplit(c, n_blocks_dimx, n_large_nodes, work_items.size(), split_smem_config); RAFT_CUDA_TRY(cudaPeekAtLastError()); } raft::update_host(h_splits, splits, work_items.size(), builder_stream); @@ -520,67 +521,67 @@ struct Builder { RAFT_CUDA_TRY(cudaPeekAtLastError()); } - size_t computeSplitHistogramSmemSize() const + struct SplitSharedMemoryConfig { + bool use_global_memory_histogram; + size_t dynamic_smem_size; + }; + + SplitSharedMemoryConfig computeSplitSharedMemoryConfig() const { + // Dynamic shared memory for the fast path: histogram, copied quantiles, and + // alignment padding for the kernel's shared-memory layout. auto shared_histogram_size = ML::checked_mul(params.max_n_bins, dataset.num_outputs, sizeof(BinT)); auto shared_quantiles_size = ML::checked_mul(params.max_n_bins, sizeof(DataT)); - auto dynamic_smem_size = + auto shared_dynamic_smem_size = ML::checked_add(shared_histogram_size, shared_quantiles_size, sizeof(int)); - - // Extra room for alignment (see alignPointer in - // computeSplitKernel) auto alignment_smem_size = ML::checked_add(sizeof(DataT), ML::checked_mul(3, sizeof(int))); - return ML::checked_add(dynamic_smem_size, alignment_smem_size); - } + shared_dynamic_smem_size = + ML::checked_add(shared_dynamic_smem_size, alignment_smem_size); - size_t computeSplitGlobalHistogramSmemSize() const - { - // shared_done only, plus conservative alignment room for alignPointer. - return ML::checked_add(sizeof(int), sizeof(int)); - } + // Dynamic shared memory for the fallback path only needs the per-block done + // flag used by the inter-block completion handshake. + auto global_dynamic_smem_size = ML::checked_add(sizeof(int), sizeof(int)); - size_t computeSplitStaticSmemSize() const - { - // computeSplitKernel also reserves static shared memory for CUB's scan temp - // storage and the per-warp split reduction scratch. + // Static shared memory is reserved by the kernel regardless of where the + // histogram lives. auto cdf_scan_smem_size = sizeof(typename cub::BlockScan::TempStorage); auto split_scratch_smem_size = ML::checked_mul(raft::ceildiv(TPB_DEFAULT, raft::WarpSize), sizeof(SplitT)); - return ML::checked_add(cdf_scan_smem_size, split_scratch_smem_size); - } + auto static_smem_size = + ML::checked_add(cdf_scan_smem_size, split_scratch_smem_size); - size_t computeSplitSmemSize() const - { - return ML::checked_add(computeSplitHistogramSmemSize(), - computeSplitStaticSmemSize()); - } - - bool shouldUseGlobalMemoryHistogram(size_t shared_histogram_dynamic_smem_size, - size_t shared_path_total_smem_size) const - { auto available_smem = size_t(handle.get_device_properties().sharedMemPerBlock); - auto global_smem = ML::checked_add(computeSplitGlobalHistogramSmemSize(), - computeSplitStaticSmemSize()); - ASSERT(available_smem >= global_smem, "Not enough shared memory for RF split bookkeeping."); - return shared_path_total_smem_size > available_smem || - shared_histogram_dynamic_smem_size > tunable_split_histogram_dynamic_smem_limit_bytes; + auto global_total_smem_size = + ML::checked_add(global_dynamic_smem_size, static_smem_size); + ASSERT(available_smem >= global_total_smem_size, + "Not enough shared memory for RF split bookkeeping."); + + // Prefer shared memory when it fits and stays small enough for good occupancy; + // otherwise use the global histogram path to avoid launch failure or slowdown. + auto shared_total_smem_size = + ML::checked_add(shared_dynamic_smem_size, static_smem_size); + bool use_global_memory_histogram = + shared_total_smem_size > available_smem || + shared_dynamic_smem_size > tunable_split_histogram_dynamic_smem_limit_bytes; + + return {use_global_memory_histogram, + use_global_memory_histogram ? global_dynamic_smem_size : shared_dynamic_smem_size}; } - void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes, size_t n_work_items) + void computeSplit(IdxT col, + size_t n_blocks_dimx, + size_t n_large_nodes, + size_t n_work_items, + const SplitSharedMemoryConfig& split_smem_config) { // if no instances to split, return if (n_blocks_dimx == 0) return; raft::common::nvtx::range fun_scope("Builder::computeSplit @builder.cuh [batched-levelalgo]"); - auto n_bins = params.max_n_bins; - auto n_classes = dataset.num_outputs; - auto shared_histogram_dynamic_smem_size = computeSplitHistogramSmemSize(); - auto shared_path_total_smem_size = computeSplitSmemSize(); - auto use_global_memory_histogram = shouldUseGlobalMemoryHistogram( - shared_histogram_dynamic_smem_size, shared_path_total_smem_size); - auto smem_size = use_global_memory_histogram ? computeSplitGlobalHistogramSmemSize() - : shared_histogram_dynamic_smem_size; + auto n_bins = params.max_n_bins; + auto n_classes = dataset.num_outputs; + auto use_global_memory_histogram = split_smem_config.use_global_memory_histogram; // if columns left to be processed lesser than `n_blks_for_cols`, shrink the blocks along dimy auto n_blocks_dimy = std::min(n_blks_for_cols, dataset.n_sampled_cols - col); dim3 grid(n_blocks_dimx, n_blocks_dimy, 1); @@ -610,7 +611,7 @@ struct Builder { seed, use_global_memory_histogram, grid, - smem_size, + split_smem_config.dynamic_smem_size, builder_stream); } From ee8383d3e2a311019a4eda9cc8d85fd1f3c9cdc3 Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Mon, 6 Jul 2026 12:23:54 +0200 Subject: [PATCH 4/4] Guard RF split histogram memset sizing --- cpp/src/decisiontree/batched-levelalgo/builder.cuh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index 567eae2768..7127f96449 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -586,8 +586,10 @@ struct Builder { auto n_blocks_dimy = std::min(n_blks_for_cols, dataset.n_sampled_cols - col); dim3 grid(n_blocks_dimx, n_blocks_dimy, 1); auto histogram_node_count = use_global_memory_histogram ? n_work_items : n_large_nodes; - size_t len_histograms = size_t(n_bins) * n_classes * n_blocks_dimy * histogram_node_count; - RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream)); + auto len_histograms = + ML::checked_mul(n_bins, n_classes, n_blocks_dimy, histogram_node_count); + auto histograms_bytes = ML::checked_mul(sizeof(BinT), len_histograms); + RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, histograms_bytes, builder_stream)); // create the objective function object ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf, params.split_criterion); // call the computeSplitKernel