From 7aceae95e9c595d376da297f1741b72efde45003 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 11 Aug 2026 14:52:43 -0700 Subject: [PATCH 01/17] Fix nvforest::handle_t when no GPU is available --- cpp/include/nvforest/handle.hpp | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/cpp/include/nvforest/handle.hpp b/cpp/include/nvforest/handle.hpp index fd1f3bb..3d51d3c 100644 --- a/cpp/include/nvforest/handle.hpp +++ b/cpp/include/nvforest/handle.hpp @@ -18,14 +18,21 @@ struct handle_t { handle_t(raft::handle_t const& raft_handle) : raft_handle_{&raft_handle} {} auto get_next_usable_stream() const { + if (!raft_handle_) { return cuda_stream{}; } return cuda_stream{raft_handle_->get_next_usable_stream().value()}; } - auto get_stream_pool_size() const { return raft_handle_->get_stream_pool_size(); } + auto get_stream_pool_size() const + { + if (!raft_handle_) { return std::size_t{}; } + return raft_handle_->get_stream_pool_size(); + } auto get_usable_stream_count() const { return std::max(get_stream_pool_size(), std::size_t{1}); } void synchronize() const { - raft_handle_->sync_stream_pool(); - raft_handle_->sync_stream(); + if (raft_handle_) { + raft_handle_->sync_stream_pool(); + raft_handle_->sync_stream(); + } } private: From 2f347784607ed4bc02fc4fbc6ab74700cea804b0 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 11 Aug 2026 22:26:33 -0700 Subject: [PATCH 02/17] Remove RMM and RAFT --- cpp/CMakeLists.txt | 9 +- cpp/cmake/thirdparty/get_raft.cmake | 78 ----------- cpp/cmake/thirdparty/get_rmm.cmake | 13 -- cpp/include/nvforest/detail/device_id/gpu.hpp | 11 +- .../nvforest/detail/device_setter/gpu.hpp | 18 ++- .../nvforest/detail/owning_buffer/gpu.hpp | 34 +++-- cpp/include/nvforest/forest_model.hpp | 125 ++++-------------- cpp/include/nvforest/handle.hpp | 50 ------- cpp/include/nvforest/treelite_importer.hpp | 1 - cpp/src/detail/device_buffer.cu | 50 +++++++ cpp/tests/CMakeLists.txt | 4 - cpp/tests/treelite_importer.cpp | 22 +-- 12 files changed, 125 insertions(+), 290 deletions(-) delete mode 100644 cpp/cmake/thirdparty/get_raft.cmake delete mode 100644 cpp/cmake/thirdparty/get_rmm.cmake delete mode 100644 cpp/include/nvforest/handle.hpp create mode 100644 cpp/src/detail/device_buffer.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 885fa54..f8d2101 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -198,10 +198,6 @@ endif() # CCCL before RMM, and RMM before RAFT include(cmake/thirdparty/get_cccl.cmake) -if(NVFOREST_ENABLE_GPU) - include(cmake/thirdparty/get_rmm.cmake) - include(cmake/thirdparty/get_raft.cmake) -endif() include(cmake/thirdparty/get_treelite.cmake) if(BUILD_NVFOREST_TESTS) @@ -252,7 +248,8 @@ if(NVFOREST_ENABLE_GPU) src/infer8.cu src/infer9.cu src/infer10.cu - src/infer11.cu) + src/infer11.cu + src/detail/device_buffer.cu) endif() target_sources( ${NVFOREST_CPP_TARGET} @@ -312,7 +309,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE) endif() if(NVFOREST_ENABLE_GPU) - list(APPEND _nvforest_cpp_public_libs rmm::rmm raft::raft CUDA::cudart_static) + list(APPEND _nvforest_cpp_public_libs CUDA::cudart_static) endif() # These are always private: diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake deleted file mode 100644 index 1265804..0000000 --- a/cpp/cmake/thirdparty/get_raft.cmake +++ /dev/null @@ -1,78 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -set(NVFOREST_MIN_VERSION_raft "${NVForest_VERSION_MAJOR}.${NVForest_VERSION_MINOR}.00") - -function(find_and_configure_raft) - set(oneValueArgs - VERSION - FORK - PINNED_TAG - EXCLUDE_FROM_ALL - USE_RAFT_STATIC - COMPILE_LIBRARY - CLONE_ON_PIN - NVTX) - cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "${rapids-cmake-checkout-tag}") - message(STATUS "NVFOREST: RAFT pinned tag found: ${PKG_PINNED_TAG}. Cloning raft locally.") - set(CPM_DOWNLOAD_raft ON) - elseif(PKG_USE_RAFT_STATIC AND (NOT CPM_raft_SOURCE)) - message(STATUS "NVFOREST: Cloning raft locally to build static libraries.") - set(CPM_DOWNLOAD_raft ON) - endif() - - # We need RAFT::distributed for MG tests - if(BUILD_NVFOREST_MG_TESTS) - string(APPEND RAFT_COMPONENTS " distributed") - endif() - - # We need to set this each time so that on subsequent calls to cmake the raft-config.cmake re-evaluates the RAFT_NVTX - # value - set(RAFT_NVTX ${PKG_NVTX}) - - message(VERBOSE "NVFOREST: raft FIND_PACKAGE_ARGUMENTS COMPONENTS ${RAFT_COMPONENTS}") - - rapids_cpm_find( - raft ${PKG_VERSION} - GLOBAL_TARGETS raft::raft - BUILD_EXPORT_SET nvforest-exports - INSTALL_EXPORT_SET nvforest-exports COMPONENTS ${RAFT_COMPONENTS} - CPM_ARGS - GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git - GIT_TAG ${PKG_PINNED_TAG} SOURCE_SUBDIR cpp - EXCLUDE_FROM_ALL ${PKG_EXCLUDE_FROM_ALL} - OPTIONS "BUILD_TESTS OFF" "BUILD_PRIMS_BENCH OFF" "BUILD_CAGRA_HNSWLIB OFF" "RAFT_COMPILE_LIBRARY OFF") - - if(raft_ADDED) - message(VERBOSE "NVFOREST: Using RAFT located in ${raft_SOURCE_DIR}") - else() - message(VERBOSE "NVFOREST: Using RAFT located in ${raft_DIR}") - endif() - -endfunction() - -# Change pinned tag here to test a commit in CI To use a different RAFT locally, set the CMake variable -# CPM_raft_SOURCE=/path/to/local/raft -find_and_configure_raft( - VERSION - ${NVFOREST_MIN_VERSION_raft} - FORK - rapidsai - PINNED_TAG - ${rapids-cmake-checkout-tag} - EXCLUDE_FROM_ALL - ${NVFOREST_EXCLUDE_RAFT_FROM_ALL} - # When PINNED_TAG above doesn't match nvforest, force local raft clone in build directory even if it's already - # installed. - CLONE_ON_PIN - ${NVFOREST_RAFT_CLONE_ON_PIN} - USE_RAFT_STATIC - ${NVFOREST_USE_RAFT_STATIC} - NVTX - ${NVTX}) diff --git a/cpp/cmake/thirdparty/get_rmm.cmake b/cpp/cmake/thirdparty/get_rmm.cmake deleted file mode 100644 index b476b09..0000000 --- a/cpp/cmake/thirdparty/get_rmm.cmake +++ /dev/null @@ -1,13 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -function(find_and_configure_rmm) - include(${rapids-cmake-dir}/cpm/rmm.cmake) - rapids_cpm_rmm(BUILD_EXPORT_SET nvforest-exports INSTALL_EXPORT_SET nvforest-exports) -endfunction() - -find_and_configure_rmm() diff --git a/cpp/include/nvforest/detail/device_id/gpu.hpp b/cpp/include/nvforest/detail/device_id/gpu.hpp index 49edfdd..10d1011 100644 --- a/cpp/include/nvforest/detail/device_id/gpu.hpp +++ b/cpp/include/nvforest/detail/device_id/gpu.hpp @@ -7,23 +7,20 @@ #include #include -#include - namespace nvforest::detail { template <> struct device_id { - using value_type = typename rmm::cuda_device_id::value_type; device_id() noexcept(false) : id_{[]() { - auto raw_id = value_type{}; + auto raw_id = int{}; cuda_check(cudaGetDevice(&raw_id)); return raw_id; }()} {}; - device_id(value_type dev_id) noexcept : id_{dev_id} {}; + device_id(int dev_id) noexcept : id_{dev_id} {}; - auto value() const noexcept { return id_.value(); } + auto value() const noexcept { return id_; } private: - rmm::cuda_device_id id_; + int id_; }; } // namespace nvforest::detail diff --git a/cpp/include/nvforest/detail/device_setter/gpu.hpp b/cpp/include/nvforest/detail/device_setter/gpu.hpp index d69a0ff..07ba9a3 100644 --- a/cpp/include/nvforest/detail/device_setter/gpu.hpp +++ b/cpp/include/nvforest/detail/device_setter/gpu.hpp @@ -8,10 +8,22 @@ #include #include -#include - #include +#include + +#define NVFOREST_CUDA_TRY_NO_THROW(call) \ + do { \ + cudaError_t const status = call; \ + if (cudaSuccess != status) { \ + printf("CUDA call='%s' at file=%s line=%d failed with %s\n", \ + #call, \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(status)); \ + } \ + } while (0) + namespace nvforest::detail { /** Struct for setting current device within a code block */ @@ -27,7 +39,7 @@ struct device_setter { cuda_check(cudaSetDevice(device.value())); } - ~device_setter() { RAFT_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); } + ~device_setter() { NVFOREST_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); } private: device_id prev_device_; diff --git a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp index ba6030f..8cca3f4 100644 --- a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp +++ b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp @@ -4,36 +4,48 @@ */ #pragma once #include -#include #include #include -#include - #include +#include +#include #include namespace nvforest::detail { + +struct owning_device_buffer_type_erased_impl; + +struct owning_device_buffer_type_erased { + owning_device_buffer_type_erased(); + owning_device_buffer_type_erased(device_id device_id, + std::size_t size, + cudaStream_t stream); + owning_device_buffer_type_erased(owning_device_buffer_type_erased&& other) noexcept; + owning_device_buffer_type_erased& operator=(owning_device_buffer_type_erased&& other) noexcept; + ~owning_device_buffer_type_erased(); + std::byte* get(); + + private: + std::unique_ptr impl_; +}; + template struct owning_buffer { // TODO(wphicks): Assess need for buffers of const T using value_type = std::remove_const_t; - owning_buffer() : data_{} {} - + owning_buffer() = default; owning_buffer(device_id device_id, std::size_t size, cudaStream_t stream) noexcept(false) - : data_{[&device_id, &size, &stream]() { - auto device_context = device_setter{device_id}; - return rmm::device_buffer{size * sizeof(value_type), rmm::cuda_stream_view{stream}}; - }()} + : data_{device_id, size * sizeof(value_type), stream} { } - auto* get() const { return reinterpret_cast(data_.data()); } + auto* get() const { return reinterpret_cast(data_.get()); } private: - mutable rmm::device_buffer data_; + mutable owning_device_buffer_type_erased data_; }; } // namespace nvforest::detail diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index e261643..f882d4c 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -6,13 +6,11 @@ #include #include #include -#include #include #include #include #include #include -#include #include #ifdef NVFOREST_ENABLE_GPU @@ -146,105 +144,30 @@ struct forest_model { { std::visit( [this, predict_type, &output, &input, &stream, &specified_chunk_size]( - auto&& concrete_forest) { - if constexpr (std::is_same_v< - typename std::remove_reference_t::io_type, - io_t>) { - concrete_forest.predict(output, input, stream, predict_type, specified_chunk_size); - } else { - throw type_error("Input type does not match model_type"); - } - }, - decision_forest_); - } - - /** - * Perform inference on given input - * - * @param[in] handle The nvforest::handle_t (wrapper for raft::handle_t - * on GPU) which will be used to provide streams for evaluation. - * @param[out] output The buffer where model output should be stored. If - * this buffer is on host while the model is on device or vice versa, - * work will be distributed across available streams to copy the data back - * to this output location. This must be of size at least ROWS x num_outputs(). - * @param[in] input The buffer containing input data. If - * this buffer is on host while the model is on device or vice versa, - * work will be distributed across available streams to copy the input data - * to the appropriate location and perform inference. - * @param[in] predict_type Type of inference to perform. Defaults to summing - * the outputs of all trees and produce an output per row. If set to - * "per_tree", we will instead output all outputs of individual trees. - * If set to "leaf_id", we will output the integer ID of the leaf node - * for each tree. - * @param[in] specified_chunk_size: Specifies the mini-batch size for - * processing. This has different meanings on CPU and GPU, but on GPU it - * corresponds to the number of rows evaluated per inference iteration - * on a single block. It can take on any power of 2 from 1 to 32, and - * runtime performance is quite sensitive to the value chosen. In general, - * larger batches benefit from higher values, but it is hard to predict the - * optimal value a priori. If omitted, a heuristic will be used to select a - * reasonable value. On CPU, this argument can generally just be omitted. - */ - template - void predict(handle_t const& handle, - buffer& output, - buffer const& input, - infer_kind predict_type = infer_kind::default_kind, - std::optional specified_chunk_size = std::nullopt) - { - std::visit( - [this, predict_type, &handle, &output, &input, &specified_chunk_size]( auto&& concrete_forest) { using model_io_t = typename std::remove_reference_t::io_type; if constexpr (std::is_same_v) { if (output.memory_type() == memory_type() && input.memory_type() == memory_type()) { - concrete_forest.predict( - output, input, handle.get_next_usable_stream(), predict_type, specified_chunk_size); + concrete_forest.predict(output, input, stream, predict_type, specified_chunk_size); } else { - auto constexpr static const MIN_CHUNKS_PER_PARTITION = std::size_t{64}; - auto constexpr static const MAX_CHUNK_SIZE = std::size_t{64}; - auto row_count = input.size() / num_features(); - auto partition_size = - std::max(detail::ceildiv(row_count, handle.get_usable_stream_count()), - specified_chunk_size.value_or(MAX_CHUNK_SIZE) * MIN_CHUNKS_PER_PARTITION); - auto partition_count = detail::ceildiv(row_count, partition_size); - for (auto i = std::size_t{}; i < partition_count; ++i) { - auto stream = handle.get_next_usable_stream(); - auto rows_in_this_partition = - std::min(partition_size, row_count - i * partition_size); - auto partition_in = buffer{}; - if (input.memory_type() != memory_type()) { - partition_in = buffer{rows_in_this_partition * num_features(), memory_type()}; - copy(partition_in, - input, - 0, - i * partition_size * num_features(), - partition_in.size(), - stream); - } else { - partition_in = buffer{input.data() + i * partition_size * num_features(), - rows_in_this_partition * num_features(), - memory_type()}; - } - auto partition_out = buffer{}; - if (output.memory_type() != memory_type()) { - partition_out = buffer{rows_in_this_partition * num_outputs(), memory_type()}; - } else { - partition_out = buffer{output.data() + i * partition_size * num_outputs(), - rows_in_this_partition * num_outputs(), - memory_type()}; - } - concrete_forest.predict( - partition_out, partition_in, stream, predict_type, specified_chunk_size); - if (output.memory_type() != memory_type()) { - copy(output, - partition_out, - i * partition_size * num_outputs(), - 0, - partition_out.size(), - stream); - } + auto buffer_in = buffer{}; + if (input.memory_type() != memory_type()) { + buffer_in = buffer{row_count * num_features(), memory_type()}; + copy(buffer_in, input, 0, 0, buffer_in.size(), stream); + } else { + buffer_in = buffer{input.data(), row_count * num_features(), memory_type()}; + } + auto buffer_out = buffer{}; + if (output.memory_type() != memory_type()) { + buffer_out = buffer{row_count * num_outputs(), memory_type()}; + } else { + buffer_out = buffer{output.data(), row_count * num_outputs(), memory_type()}; + } + concrete_forest.predict( + buffer_out, buffer_in, stream, predict_type, specified_chunk_size); + if (output.memory_type() != memory_type()) { + copy(output, buffer_out, 0, 0, buffer_out.size(), stream); } } } else { @@ -257,8 +180,6 @@ struct forest_model { /** * Perform inference on given input * - * @param[in] handle The nvforest::handle_t (wrapper for raft::handle_t - * on GPU) which will be used to provide streams for evaluation. * @param[out] output Pointer to the memory location where output should end * up * @param[in] input Pointer to the input data @@ -266,6 +187,10 @@ struct forest_model { * @param[in] out_mem_type The memory type (device/host) of the output * buffer * @param[in] in_mem_type The memory type (device/host) of the input buffer + * @param[in] stream A nvforest::cuda_stream, which (on GPU-enabled builds) is + * a transparent wrapper for the cudaStream_t or (on CPU-only builds) a + * CUDA-free placeholder object. If this argument is cudaStream_t, ensure + * that this stream is associated with the same GPU device as the model object. * @param[in] predict_type Type of inference to perform. Defaults to summing * the outputs of all trees and produce an output per row. If set to * "per_tree", we will instead output all outputs of individual trees. @@ -281,12 +206,12 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(handle_t const& handle, - io_t* output, + void predict(io_t* output, io_t* input, std::size_t num_rows, device_type out_mem_type, device_type in_mem_type, + cuda_stream stream = cuda_stream{}, infer_kind predict_type = infer_kind::default_kind, std::optional specified_chunk_size = std::nullopt) { @@ -302,7 +227,7 @@ struct forest_model { } auto out_buffer = buffer{output, num_rows * num_outputs(), out_mem_type, current_device_id}; auto in_buffer = buffer{input, num_rows * num_features(), in_mem_type, current_device_id}; - predict(handle, out_buffer, in_buffer, predict_type, specified_chunk_size); + predict(out_buffer, in_buffer, stream, predict_type, specified_chunk_size); } private: diff --git a/cpp/include/nvforest/handle.hpp b/cpp/include/nvforest/handle.hpp deleted file mode 100644 index 3d51d3c..0000000 --- a/cpp/include/nvforest/handle.hpp +++ /dev/null @@ -1,50 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once -#include - -#include -#include -#ifdef NVFOREST_ENABLE_GPU -#include -#endif - -namespace nvforest { -#ifdef NVFOREST_ENABLE_GPU -struct handle_t { - handle_t(raft::handle_t const* handle_ptr = nullptr) : raft_handle_{handle_ptr} {} - handle_t(raft::handle_t const& raft_handle) : raft_handle_{&raft_handle} {} - auto get_next_usable_stream() const - { - if (!raft_handle_) { return cuda_stream{}; } - return cuda_stream{raft_handle_->get_next_usable_stream().value()}; - } - auto get_stream_pool_size() const - { - if (!raft_handle_) { return std::size_t{}; } - return raft_handle_->get_stream_pool_size(); - } - auto get_usable_stream_count() const { return std::max(get_stream_pool_size(), std::size_t{1}); } - void synchronize() const - { - if (raft_handle_) { - raft_handle_->sync_stream_pool(); - raft_handle_->sync_stream(); - } - } - - private: - // Have to store a pointer because handle is not movable - raft::handle_t const* raft_handle_; -}; -#else -struct handle_t { - auto get_next_usable_stream() const { return cuda_stream{}; } - auto get_stream_pool_size() const { return std::size_t{}; } - auto get_usable_stream_count() const { return std::max(get_stream_pool_size(), std::size_t{1}); } - void synchronize() const {} -}; -#endif -} // namespace nvforest diff --git a/cpp/include/nvforest/treelite_importer.hpp b/cpp/include/nvforest/treelite_importer.hpp index 5402b88..b504403 100644 --- a/cpp/include/nvforest/treelite_importer.hpp +++ b/cpp/include/nvforest/treelite_importer.hpp @@ -18,7 +18,6 @@ #include #include -#include #include #include diff --git a/cpp/src/detail/device_buffer.cu b/cpp/src/detail/device_buffer.cu new file mode 100644 index 0000000..a9fdf21 --- /dev/null +++ b/cpp/src/detail/device_buffer.cu @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include +#include +#include + +#include +#include +#include + +#include +#include + +namespace nvforest::detail { + +struct owning_device_buffer_type_erased_impl { + owning_device_buffer_type_erased_impl(int device_id, std::size_t size, cudaStream_t stream) + : buffer_{[&stream, device_id, size]() { + auto device = cuda::device_ref{device_id}; + auto mr = cuda::device_default_memory_pool(device); + return cuda::make_buffer(cuda::stream_ref{stream}, mr, size, cuda::no_init); + }()} + { + } + + std::byte* get() { return buffer_.data(); } + + cuda::device_buffer buffer_; +}; + +owning_device_buffer_type_erased::owning_device_buffer_type_erased() : impl_{nullptr} {} + +owning_device_buffer_type_erased::owning_device_buffer_type_erased( + device_id device_id, std::size_t size, cudaStream_t stream) +{ + auto device_context = device_setter{device_id}; + impl_ = std::make_unique(device_id.value(), size, stream); +} + +owning_device_buffer_type_erased::owning_device_buffer_type_erased( + owning_device_buffer_type_erased&& other) noexcept = default; +owning_device_buffer_type_erased& owning_device_buffer_type_erased::operator=( + owning_device_buffer_type_erased&& other) noexcept = default; +owning_device_buffer_type_erased::~owning_device_buffer_type_erased() = default; + +std::byte* owning_device_buffer_type_erased::get() { return impl_->get(); } + +} // namespace nvforest::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 96f5297..d4aeed8 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -51,10 +51,6 @@ function(ConfigureTest) Threads::Threads ${TREELITE_LIBS} $) - if(NVFOREST_ENABLE_GPU) - target_link_libraries(${_NVFOREST_TEST_NAME} PRIVATE rmm::rmm raft::raft) - endif() - target_compile_options(${_NVFOREST_TEST_NAME} PRIVATE "$<$:${NVFOREST_CXX_FLAGS}>" "$<$:${NVFOREST_CUDA_FLAGS}>") diff --git a/cpp/tests/treelite_importer.cpp b/cpp/tests/treelite_importer.cpp index 37d3d9c..1b7f089 100644 --- a/cpp/tests/treelite_importer.cpp +++ b/cpp/tests/treelite_importer.cpp @@ -3,8 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include -#include #include #include #include @@ -382,21 +382,15 @@ TEST(TreeliteImporter, DegenerateTree) auto nvforest_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); ASSERT_FALSE(nvforest_model.has_vector_leaves()); -#ifdef NVFOREST_ENABLE_GPU - auto raft_handle = raft::handle_t{}; - auto handle = nvforest::handle_t{raft_handle}; -#else - auto handle = nvforest::handle_t{}; -#endif auto X = std::vector{0.0}; auto preds = std::vector(1, 0.0); auto expected_preds = std::vector{1.0}; - nvforest_model.predict(handle, - preds.data(), + nvforest_model.predict(preds.data(), X.data(), 1, nvforest::device_type::cpu, nvforest::device_type::cpu, + nvforest::cuda_stream{}, nvforest::infer_kind::default_kind, 1); ASSERT_EQ(preds, expected_preds); @@ -408,21 +402,15 @@ TEST(TreeliteImporter, DegenerateTreeWithVectorLeaf) auto nvforest_model = import_from_treelite_model(*tl_model, tree_layout::breadth_first); ASSERT_TRUE(nvforest_model.has_vector_leaves()); -#ifdef NVFOREST_ENABLE_GPU - auto raft_handle = raft::handle_t{}; - auto handle = nvforest::handle_t{raft_handle}; -#else - auto handle = nvforest::handle_t{}; -#endif auto X = std::vector{0.0}; auto preds = std::vector(2, 0.0); auto expected_preds = std::vector{0.5, 0.5}; - nvforest_model.predict(handle, - preds.data(), + nvforest_model.predict(preds.data(), X.data(), 1, nvforest::device_type::cpu, nvforest::device_type::cpu, + nvforest::cuda_stream{}, nvforest::infer_kind::default_kind, 1); ASSERT_EQ(preds, expected_preds); From d53522da4713f076b6376f1313e2d575cc8a7abd Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 00:24:50 -0700 Subject: [PATCH 03/17] Update Python pkg + pytests --- ci/build_wheel_libnvforest.sh | 4 +- ci/build_wheel_nvforest.sh | 2 - ci/release/update-version.sh | 5 - .../all_cuda-129_arch-aarch64.yaml | 4 +- .../all_cuda-129_arch-x86_64.yaml | 4 +- .../all_cuda-133_arch-aarch64.yaml | 4 +- .../all_cuda-133_arch-x86_64.yaml | 4 +- .../clang_tidy_cuda-129_arch-x86_64.yaml | 2 - .../clang_tidy_cuda-133_arch-x86_64.yaml | 2 - .../cpp_all_cuda-129_arch-x86_64.yaml | 2 - .../cpp_all_cuda-133_arch-x86_64.yaml | 2 - conda/recipes/libnvforest/recipe.yaml | 6 -- conda/recipes/nvforest/recipe.yaml | 4 +- dependencies.yaml | 102 ++---------------- python/libnvforest/libnvforest/load.py | 4 - python/libnvforest/pyproject.toml | 4 - python/nvforest/nvforest/__init__.py | 4 +- python/nvforest/nvforest/_base.py | 22 ++-- python/nvforest/nvforest/_factory.py | 44 ++++---- python/nvforest/nvforest/_forest_inference.py | 42 ++++---- python/nvforest/nvforest/_handle.py | 9 -- .../nvforest/nvforest/detail/cuda_stream.pxd | 9 +- .../nvforest/detail/forest_inference.pyx | 54 ++++++---- python/nvforest/nvforest/detail/handle.pxd | 17 --- python/nvforest/pyproject.toml | 5 +- python/nvforest/tests/test_nvforest.py | 31 ++++++ 26 files changed, 138 insertions(+), 254 deletions(-) delete mode 100644 python/nvforest/nvforest/_handle.py delete mode 100644 python/nvforest/nvforest/detail/handle.pxd diff --git a/ci/build_wheel_libnvforest.sh b/ci/build_wheel_libnvforest.sh index b6b8db8..d03ca41 100755 --- a/ci/build_wheel_libnvforest.sh +++ b/ci/build_wheel_libnvforest.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -29,7 +29,6 @@ rapids-pip-retry install \ export PIP_NO_BUILD_ISOLATION=0 EXCLUDE_ARGS=( - --exclude "libraft.so" --exclude "libcublas.so.*" --exclude "libcublasLt.so.*" --exclude "libcufft.so.*" @@ -38,7 +37,6 @@ EXCLUDE_ARGS=( --exclude "libcusparse.so.*" --exclude "libnvJitLink.so.*" --exclude "librapids_logger.so" - --exclude "librmm.so" ) export SKBUILD_CMAKE_ARGS="-DDISABLE_DEPRECATION_WARNINGS=ON" diff --git a/ci/build_wheel_nvforest.sh b/ci/build_wheel_nvforest.sh index 36c7e8a..e86fe3d 100755 --- a/ci/build_wheel_nvforest.sh +++ b/ci/build_wheel_nvforest.sh @@ -21,7 +21,6 @@ echo "libnvforest-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBNVFOREST_WHEELH EXCLUDE_ARGS=( --exclude "libnvforest.so" - --exclude "libraft.so" --exclude "libcublas.so.*" --exclude "libcublasLt.so.*" --exclude "libcufft.so.*" @@ -30,7 +29,6 @@ EXCLUDE_ARGS=( --exclude "libcusparse.so.*" --exclude "libnvJitLink.so.*" --exclude "librapids_logger.so" - --exclude "librmm.so" --exclude "libtreelite.so" ) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index bd552e7..42e57b9 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -109,11 +109,6 @@ DEPENDENCIES=( nvforest libnvforest libnvforest-tests - libraft - libraft-headers - librmm - pylibraft - rmm ) for DEP in "${DEPENDENCIES[@]}"; do for FILE in dependencies.yaml conda/environments/*.yaml; do diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 6dd75f5..6f43053 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=12.9.2,<13.0 +- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api @@ -21,15 +22,12 @@ dependencies: - doxygen - gcc_linux-aarch64=14.* - graphviz -- libraft==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc - packaging - pandas - pre-commit -- pylibraft==26.10.*,>=0.0.0a0 - pytest - pytest-cov - pytest-xdist diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index d9e43d6..69c5031 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=12.9.2,<13.0 +- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api @@ -21,15 +22,12 @@ dependencies: - doxygen - gcc_linux-64=14.* - graphviz -- libraft==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc - packaging - pandas - pre-commit -- pylibraft==26.10.*,>=0.0.0a0 - pytest - pytest-cov - pytest-xdist diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index fdbe67e..49e3f6d 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=13.0.1,<14.0 +- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api @@ -21,15 +22,12 @@ dependencies: - doxygen - gcc_linux-aarch64=14.* - graphviz -- libraft==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc - packaging - pandas - pre-commit -- pylibraft==26.10.*,>=0.0.0a0 - pytest - pytest-cov - pytest-xdist diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index eb3d2c2..ef872a2 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -11,6 +11,7 @@ dependencies: - click - cmake>=4.0 - cuda-bindings>=13.0.1,<14.0 +- cuda-core>=1.0.0,<2.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-profiler-api @@ -21,15 +22,12 @@ dependencies: - doxygen - gcc_linux-64=14.* - graphviz -- libraft==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc - packaging - pandas - pre-commit -- pylibraft==26.10.*,>=0.0.0a0 - pytest - pytest-cov - pytest-xdist diff --git a/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml b/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml index a7c09a5..8aff852 100644 --- a/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml +++ b/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml @@ -15,8 +15,6 @@ dependencies: - cuda-version=12.9 - cxx-compiler - gcc_linux-64=14.* -- libraft-headers==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - llvm-openmp==15.0.7 - ninja - sysroot_linux-64==2.28 diff --git a/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml b/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml index ce26c36..1fc03b9 100644 --- a/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml +++ b/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml @@ -15,8 +15,6 @@ dependencies: - cuda-version=13.3 - cxx-compiler - gcc_linux-64=14.* -- libraft-headers==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - llvm-openmp==15.0.7 - ninja - sysroot_linux-64==2.28 diff --git a/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml b/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml index 0923646..b22a842 100644 --- a/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml @@ -13,8 +13,6 @@ dependencies: - cuda-version=12.9 - cxx-compiler - gcc_linux-64=14.* -- libraft-headers==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - ninja - sysroot_linux-64==2.28 name: cpp_all_cuda-129_arch-x86_64 diff --git a/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml b/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml index 744248b..c9774be 100644 --- a/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml @@ -13,8 +13,6 @@ dependencies: - cuda-version=13.3 - cxx-compiler - gcc_linux-64=14.* -- libraft-headers==26.10.*,>=0.0.0a0 -- librmm==26.10.*,>=0.0.0a0 - ninja - sysroot_linux-64==2.28 name: cpp_all_cuda-133_arch-x86_64 diff --git a/conda/recipes/libnvforest/recipe.yaml b/conda/recipes/libnvforest/recipe.yaml index e58856d..fe69c81 100644 --- a/conda/recipes/libnvforest/recipe.yaml +++ b/conda/recipes/libnvforest/recipe.yaml @@ -72,8 +72,6 @@ cache: - ${{ stdlib("c") }} host: - cuda-version =${{ cuda_version }} - - libraft-headers =${{ minor_version }} - - librmm =${{ minor_version }} - rapids-logger =0.3 - treelite ${{ treelite_version }} - cuda-cudart-dev @@ -100,19 +98,16 @@ outputs: host: - cuda-version =${{ cuda_version }} - rapids-logger =0.3 - - librmm =${{ minor_version }} - treelite ${{ treelite_version }} - cuda-cudart-dev run: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - - librmm =${{ minor_version }} - rapids-logger =0.3 - treelite ${{ treelite_version }} ignore_run_exports: by_name: - cuda-cudart - cuda-version - - librmm - libcusolver about: homepage: ${{ load_from_file("python/libnvforest/pyproject.toml").project.urls.Homepage }} @@ -143,7 +138,6 @@ outputs: by_name: - cuda-cudart - cuda-version - - librmm about: homepage: ${{ load_from_file("python/libnvforest/pyproject.toml").project.urls.Homepage }} license: ${{ load_from_file("python/libnvforest/pyproject.toml").project.license }} diff --git a/conda/recipes/nvforest/recipe.yaml b/conda/recipes/nvforest/recipe.yaml index 5a7faa4..085ebbc 100644 --- a/conda/recipes/nvforest/recipe.yaml +++ b/conda/recipes/nvforest/recipe.yaml @@ -80,7 +80,7 @@ requirements: - cython >=3.0.0 - libnvforest =${{ version }} - pip - - pylibraft =${{ minor_version }} + - cuda-core >=1.0.0,<2.0.0 - python =${{ py_abi_min }} - python-abi3 ${{ py_abi_min }}.* - rapids-build-backend >=0.4.0,<0.5.0.dev0 @@ -96,7 +96,7 @@ requirements: - libnvforest =${{ version }} - numpy >=2.0,<3.0a0 - scikit-learn >=1.4 - - pylibraft =${{ minor_version }} + - cuda-core >=1.0.0,<2.0.0 - python - treelite ${{ treelite_version }} - cuda-cudart diff --git a/dependencies.yaml b/dependencies.yaml index ddcbbe2..0f2ed6c 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,10 +13,8 @@ files: - cuda - cuda_version - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - - depends_on_libraft - - depends_on_librmm - - depends_on_pylibraft - depends_on_rapids_logger - develop - docs @@ -37,10 +35,8 @@ files: - cuda - cuda_version - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - - depends_on_libraft - - depends_on_librmm - - depends_on_pylibraft - depends_on_rapids_logger - py_build_nvforest - py_run_nvforest @@ -56,8 +52,6 @@ files: - common_build - cuda - cuda_version - - depends_on_libraft_headers - - depends_on_librmm checks: output: none includes: @@ -73,8 +67,6 @@ files: - common_build - cuda - cuda_version - - depends_on_libraft_headers - - depends_on_librmm docs: output: none includes: @@ -116,9 +108,6 @@ files: - common_build - depends_on_cuda_python - depends_on_libnvforest - - depends_on_libraft - - depends_on_librmm - - depends_on_pylibraft - py_build_nvforest py_run_nvforest: output: pyproject @@ -127,9 +116,9 @@ files: table: project includes: - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libnvforest - - depends_on_pylibraft - py_run_nvforest py_test_nvforest: output: pyproject @@ -155,8 +144,6 @@ files: key: requires includes: - common_build - - depends_on_libraft - - depends_on_librmm - depends_on_rapids_logger py_run_libnvforest: output: pyproject @@ -164,8 +151,6 @@ files: extras: table: project includes: - - depends_on_libraft - - depends_on_librmm - depends_on_rapids_logger channels: - rapidsai-nightly @@ -455,6 +440,11 @@ dependencies: - matrix: packages: - xgboost>=3.3.0; python_version >= '3.12' + depends_on_cuda_core: + common: + - output_types: [conda, requirements, pyproject] + packages: + - cuda-core>=1.0.0,<2.0.0 depends_on_cuda_python: specific: - output_types: [conda, requirements, pyproject] @@ -534,82 +524,6 @@ dependencies: packages: - libnvforest-cu13==26.10.*,>=0.0.0a0 - {matrix: null, packages: [*libnvforest_unsuffixed]} - depends_on_libraft: - common: - - output_types: conda - packages: - - &libraft_unsuffixed libraft==26.10.*,>=0.0.0a0 - - output_types: requirements - packages: - # pip recognizes the index as a global option for the requirements.txt file - - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple - specific: - - output_types: [requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - cuda_suffixed: "true" - packages: - - libraft-cu12==26.10.*,>=0.0.0a0 - - matrix: - cuda: "13.*" - cuda_suffixed: "true" - packages: - - libraft-cu13==26.10.*,>=0.0.0a0 - - {matrix: null, packages: [*libraft_unsuffixed]} - depends_on_libraft_headers: - common: - - output_types: conda - packages: - - libraft-headers==26.10.*,>=0.0.0a0 - depends_on_librmm: - common: - - output_types: conda - packages: - - &librmm_unsuffixed librmm==26.10.*,>=0.0.0a0 - - output_types: requirements - packages: - # pip recognizes the index as a global option for the requirements.txt file - - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple - specific: - - output_types: [requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - cuda_suffixed: "true" - packages: - - librmm-cu12==26.10.*,>=0.0.0a0 - - matrix: - cuda: "13.*" - cuda_suffixed: "true" - packages: - - librmm-cu13==26.10.*,>=0.0.0a0 - - {matrix: null, packages: [*librmm_unsuffixed]} - depends_on_pylibraft: - common: - - output_types: conda - packages: - - &pylibraft_unsuffixed pylibraft==26.10.*,>=0.0.0a0 - - output_types: requirements - packages: - # pip recognizes the index as a global option for the requirements.txt file - - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple - specific: - - output_types: [requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - cuda_suffixed: "true" - packages: - - pylibraft-cu12==26.10.*,>=0.0.0a0 - - matrix: - cuda: "13.*" - cuda_suffixed: "true" - packages: - - pylibraft-cu13==26.10.*,>=0.0.0a0 - - matrix: - packages: - - *pylibraft_unsuffixed depends_on_rapids_logger: common: - output_types: [conda, requirements, pyproject] diff --git a/python/libnvforest/libnvforest/load.py b/python/libnvforest/libnvforest/load.py index 462d1fc..88fede2 100644 --- a/python/libnvforest/libnvforest/load.py +++ b/python/libnvforest/libnvforest/load.py @@ -37,13 +37,9 @@ def load_library(): """Dynamically load libnvforest.so and its dependencies""" try: # These libraries must all be loaded before libnvforest - import libraft - import librmm import rapids_logger rapids_logger.load_library() - librmm.load_library() - libraft.load_library() except ModuleNotFoundError: # These runtime dependencies might be satisfied by conda packages (which do not # have any Python modules) instead of wheels. In that situation, assume that diff --git a/python/libnvforest/pyproject.toml b/python/libnvforest/pyproject.toml index 5765bd2..ac4cf66 100644 --- a/python/libnvforest/pyproject.toml +++ b/python/libnvforest/pyproject.toml @@ -25,8 +25,6 @@ classifiers = [ "Environment :: GPU :: NVIDIA CUDA", ] dependencies = [ - "libraft==26.10.*,>=0.0.0a0", - "librmm==26.10.*,>=0.0.0a0", "rapids-logger==0.3.*", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. @@ -64,8 +62,6 @@ dependencies-file = "../../dependencies.yaml" matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true" requires = [ "cmake>=4.0", - "libraft==26.10.*,>=0.0.0a0", - "librmm==26.10.*,>=0.0.0a0", "ninja", "rapids-logger==0.3.*", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. diff --git a/python/nvforest/nvforest/__init__.py b/python/nvforest/nvforest/__init__.py index 4fc9b15..3fe9542 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -24,7 +24,6 @@ GPUForestInferenceClassifier, GPUForestInferenceRegressor, ) -from nvforest._handle import Handle from nvforest._version import __git_commit__, __version__ __all__ = [ @@ -32,7 +31,6 @@ "CPUForestInferenceRegressor", "GPUForestInferenceClassifier", "GPUForestInferenceRegressor", - "Handle", "load_model", "load_from_sklearn", "load_from_treelite_model", diff --git a/python/nvforest/nvforest/_base.py b/python/nvforest/nvforest/_base.py index 2e635d3..1d983d3 100644 --- a/python/nvforest/nvforest/_base.py +++ b/python/nvforest/nvforest/_base.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod @@ -34,9 +34,7 @@ def predict_per_tree( to the correct datatype before inference. If this input is in a memory location that is inaccessible to the current device type (as set with the 'device' parameter in the constructor), - it will be copied to the correct location. This copy will be - distributed across as many CUDA streams as are available - in the stream pool of the model's RAFT handle. + it will be copied to the correct location on the model's CUDA stream. chunk_size : The number of rows to simultaneously process in one iteration of the inference algorithm. Batches are further broken down into @@ -72,9 +70,7 @@ def apply( to the correct datatype before inference. If this input is in a memory location that is inaccessible to the current device type (as set with the 'device' parameter in the constructor), - it will be copied to the correct location. This copy will be - distributed across as many CUDA streams as are available - in the stream pool of the model's RAFT handle. + it will be copied to the correct location on the model's CUDA stream. chunk_size : The number of rows to simultaneously process in one iteration of the inference algorithm. Batches are further broken down into @@ -226,9 +222,7 @@ def predict( to the correct datatype before inference. If this input is in a memory location that is inaccessible to the current device type (as set with the 'device' parameter in the constructor), - it will be copied to the correct location. This copy will be - distributed across as many CUDA streams as are available - in the stream pool of the model's RAFT handle. + it will be copied to the correct location on the model's CUDA stream. chunk_size : The number of rows to simultaneously process in one iteration of the inference algorithm. Batches are further broken down into @@ -270,9 +264,7 @@ def predict_proba( to the correct datatype before inference. If this input is in a memory location that is inaccessible to the current device type (as set with the 'device' parameter in the constructor), - it will be copied to the correct location. This copy will be - distributed across as many CUDA streams as are available - in the stream pool of the model's RAFT handle. + it will be copied to the correct location on the model's CUDA stream. chunk_size : The number of rows to simultaneously process in one iteration of the inference algorithm. Batches are further broken down into @@ -314,9 +306,7 @@ def predict( to the correct datatype before inference. If this input is in a memory location that is inaccessible to the current device type (as set with the 'device' parameter in the constructor), - it will be copied to the correct location. This copy will be - distributed across as many CUDA streams as are available - in the stream pool of the model's RAFT handle. + it will be copied to the correct location on the model's CUDA stream. chunk_size : The number of rows to simultaneously process in one iteration of the inference algorithm. Batches are further broken down into diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index cc0700a..37da9a7 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -7,6 +7,7 @@ from typing import Any, Optional, Union import treelite +from cuda.core import Stream from nvforest._base import ForestInference from nvforest._forest_inference import ( @@ -17,7 +18,6 @@ infer_device, infer_is_classifier, ) -from nvforest._handle import Handle def get_forest_inference_class(device, is_classifier) -> type: @@ -39,7 +39,7 @@ def make_forest_inference_object( treelite_model: treelite.Model, device: str, device_id: Optional[int], - handle: Optional[Handle], + stream: Optional[Stream], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -50,7 +50,7 @@ def make_forest_inference_object( kwargs = dict( treelite_model=treelite_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -72,7 +72,7 @@ def load_model( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, ) -> ForestInference: """Load a model into nvForest from a serialized model file. @@ -113,10 +113,10 @@ def load_model( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - handle : nvforest.Handle or None - For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. If not given, a new - handle will be constructed. + stream : cuda.core.Stream or None, default=None + For GPU execution, the CUDA stream to use during model loading and + inference. If not given, a new stream will be created. For CPU + execution, this value is ignored. """ model_path = pathlib.Path(model_file) if not model_path.exists(): @@ -157,7 +157,7 @@ def load_model( treelite_model=tl_model, device=device, device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -174,7 +174,7 @@ def load_from_sklearn( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, ) -> ForestInference: """Load a Scikit-Learn forest model to nvForest @@ -208,10 +208,10 @@ def load_from_sklearn( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - handle : nvforest.Handle or None - For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. If not given, a new - handle will be constructed. + stream : cuda.core.Stream or None, default=None + For GPU execution, the CUDA stream to use during model loading and + inference. If not given, a new stream will be created. For CPU + execution, this value is ignored. """ tl_model = treelite.sklearn.import_model(skl_model) @@ -219,7 +219,7 @@ def load_from_sklearn( treelite_model=tl_model, device=device, device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -236,7 +236,7 @@ def load_from_treelite_model( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, ) -> ForestInference: """Load a Treelite forest model to nvForest @@ -270,16 +270,16 @@ def load_from_treelite_model( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - handle : nvforest.Handle or None - For GPU execution, the nvForest handle containing the stream or stream - pool to use during loading and inference. If not given, a new - handle will be constructed. + stream : cuda.core.Stream or None, default=None + For GPU execution, the CUDA stream to use during model loading and + inference. If not given, a new stream will be created. For CPU + execution, this value is ignored. """ return make_forest_inference_object( treelite_model=tl_model, device=device, device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index f6bbbd3..e312776 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -16,9 +16,9 @@ import numpy as np import treelite from cuda.bindings import runtime +from cuda.core import Stream from nvforest._base import ForestInferenceClassifier, ForestInferenceRegressor -from nvforest._handle import Handle from nvforest._typing import DataType from nvforest.detail.forest_inference import ForestInferenceImpl @@ -123,7 +123,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[Stream], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -200,7 +200,7 @@ def _optimize( else: test_instances[layout] = type(self)._create_with_layout( treelite_model_bytes=self.forest.treelite_model_bytes, - handle=self.forest.handle, + stream=self.forest.stream, layout=layout, default_chunk_size=None, align_bytes=self.forest.align_bytes, @@ -243,7 +243,7 @@ def _optimize( # Return a new instance with optimal settings return type(self)._create_with_layout( treelite_model_bytes=self.forest.treelite_model_bytes, - handle=self.forest.handle, + stream=self.forest.stream, layout=optimal_layout, default_chunk_size=optimal_chunk_size, align_bytes=self.forest.align_bytes, @@ -260,7 +260,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -272,7 +272,7 @@ def __init__( treelite_model=treelite_model, device="cpu", device_id=-1, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -284,7 +284,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[Stream], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -296,7 +296,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -397,7 +397,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -409,7 +409,7 @@ def __init__( treelite_model=treelite_model, device="cpu", device_id=-1, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -421,7 +421,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[Stream], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -433,7 +433,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -526,7 +526,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -539,7 +539,7 @@ def __init__( treelite_model=treelite_model, device="gpu", device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -551,7 +551,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[Stream], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -563,7 +563,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -665,7 +665,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -678,7 +678,7 @@ def __init__( treelite_model=treelite_model, device="gpu", device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -690,7 +690,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[Stream], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -702,7 +702,7 @@ def _create_with_layout( tl_model = treelite.Model.deserialize_bytes(treelite_model_bytes) return cls( treelite_model=tl_model, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, diff --git a/python/nvforest/nvforest/_handle.py b/python/nvforest/nvforest/_handle.py deleted file mode 100644 index 347b748..0000000 --- a/python/nvforest/nvforest/_handle.py +++ /dev/null @@ -1,9 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# - -from pylibraft.common.handle import Handle as RaftHandle - -# For now, nvforest.handle.Handle is an alias of pylibraft.common.handle.Handle -Handle = RaftHandle diff --git a/python/nvforest/nvforest/detail/cuda_stream.pxd b/python/nvforest/nvforest/detail/cuda_stream.pxd index 690c0c4..807030b 100644 --- a/python/nvforest/nvforest/detail/cuda_stream.pxd +++ b/python/nvforest/nvforest/detail/cuda_stream.pxd @@ -2,6 +2,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -cdef extern from "nvforest/cuda_stream.hpp" namespace "nvforest" nogil: - cdef cppclass cuda_stream: +cdef extern from "cuda_runtime_api.h" nogil: + ctypedef struct CUstream_st: pass + ctypedef CUstream_st* cudaStream_t + + +cdef extern from "nvforest/cuda_stream.hpp" namespace "nvforest" nogil: + ctypedef cudaStream_t cuda_stream diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 3d6eb3b..ced0918 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -7,19 +7,17 @@ from typing import Optional, Union import numpy as np import treelite +from cuda.core import Device, Stream -from nvforest._handle import Handle from nvforest._typing import DataType from nvforest.detail.treelite import safe_treelite_call from libc.stdint cimport uint32_t, uintptr_t from libcpp cimport bool from libcpp.optional cimport nullopt, optional -from pylibraft.common.handle cimport handle_t as raft_handle_t from nvforest.detail.cuda_stream cimport cuda_stream as nvforest_stream_t from nvforest.detail.device_type cimport device_type as nvforest_device_t -from nvforest.detail.handle cimport handle_t as nvforest_handle_t from nvforest.detail.infer_kind cimport infer_kind from nvforest.detail.postprocessing cimport element_op, row_op from nvforest.detail.tree_layout cimport tree_layout as nvforest_tree_layout @@ -33,12 +31,12 @@ from nvforest.detail.treelite cimport ( cdef extern from "nvforest/forest_model.hpp" namespace "nvforest" nogil: cdef cppclass forest_model: void predict[io_t]( - const nvforest_handle_t&, io_t*, io_t*, size_t, nvforest_device_t, nvforest_device_t, + nvforest_stream_t, infer_kind, optional[uint32_t] ) except + @@ -65,13 +63,13 @@ cdef extern from "nvforest/treelite_importer.hpp" namespace "nvforest" nogil: cdef class ForestInference_impl(): cdef forest_model model - cdef nvforest_handle_t nvforest_handle - cdef object raft_handle + cdef nvforest_stream_t stream_handle + cdef object stream cdef object device def __cinit__( self, - raft_handle: object, + stream: object, tl_model_bytes: Union[bytes, bytearray], *, layout: str = "depth_first", @@ -80,12 +78,13 @@ cdef class ForestInference_impl(): device: str = "cpu", device_id: Optional[int] = None, ): - # Store reference to RAFT handle to control lifetime, since - # nvforest_handle keeps a pointer to it - self.raft_handle = raft_handle - self.nvforest_handle = nvforest_handle_t( - self.raft_handle.getHandle() - ) + cdef uintptr_t stream_ptr = 0 + if stream is not None: + if not isinstance(stream, Stream): + raise TypeError("stream must be a cuda.core.Stream or None") + stream_ptr = int(stream.handle) + self.stream = stream + self.stream_handle = stream_ptr cdef optional[bool] use_double_precision_c cdef bool use_double_precision_bool @@ -130,7 +129,7 @@ cdef class ForestInference_impl(): use_double_precision_c, dev_type, device_id, - self.nvforest_handle.get_next_usable_stream() + self.stream_handle ) safe_treelite_call( @@ -240,29 +239,29 @@ cdef class ForestInference_impl(): if model_dtype == np.float32: self.model.predict[float]( - self.nvforest_handle, out_ptr, in_ptr, n_rows, out_dev, in_dev, + self.stream_handle, infer_type_enum, chunk_specification ) else: self.model.predict[double]( - self.nvforest_handle, out_ptr, in_ptr, n_rows, out_dev, in_dev, + self.stream_handle, infer_type_enum, chunk_specification ) if self.device == "gpu": - self.nvforest_handle.synchronize() + self.stream.sync() return preds @@ -273,7 +272,7 @@ class ForestInferenceImpl: treelite_model: treelite.Model, device: str, device_id: int, - handle: Optional[Handle] = None, + stream: Optional[Stream] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -281,7 +280,22 @@ class ForestInferenceImpl: ): # Assumption: The caller needs to pass in correct (device, device_id) pair # This function will not contain any logic for auto-detecting device. - self.handle = Handle() if handle is None else handle + if stream is not None and not isinstance(stream, Stream): + raise TypeError("stream must be a cuda.core.Stream or None") + if device == "gpu" and stream is None: + previous_device = Device() + try: + cuda_device = Device(device_id) + cuda_device.set_current() + stream = cuda_device.create_stream() + finally: + previous_device.set_current() + if device == "gpu" and stream.device.device_id != device_id: + raise ValueError( + f"stream is associated with device {stream.device.device_id}, " + f"but device_id is {device_id}" + ) + self.stream = stream self._layout = layout self.precision = precision self.default_chunk_size = default_chunk_size @@ -309,7 +323,7 @@ class ForestInferenceImpl: self._treelite_model_bytes = treelite_model.serialize_bytes() self.impl = ForestInference_impl( - self.handle, + self.stream, self._treelite_model_bytes, layout=self._layout, align_bytes=self.align_bytes, diff --git a/python/nvforest/nvforest/detail/handle.pxd b/python/nvforest/nvforest/detail/handle.pxd deleted file mode 100644 index b80ad6b..0000000 --- a/python/nvforest/nvforest/detail/handle.pxd +++ /dev/null @@ -1,17 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# - -from pylibraft.common.handle cimport handle_t as raft_handle_t - -from nvforest.detail.cuda_stream cimport cuda_stream as nvforest_stream_t - - -cdef extern from "nvforest/handle.hpp" namespace "nvforest" nogil: - cdef cppclass handle_t: - handle_t() except + - handle_t(const raft_handle_t* handle_ptr) except + - handle_t(const raft_handle_t& handle) except + - nvforest_stream_t get_next_usable_stream() except + - void synchronize() except+ diff --git a/python/nvforest/pyproject.toml b/python/nvforest/pyproject.toml index 437fb9d..2f0fd1a 100644 --- a/python/nvforest/pyproject.toml +++ b/python/nvforest/pyproject.toml @@ -28,10 +28,10 @@ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ "cuda-bindings>=13.0.1,<14.0", + "cuda-core>=1.0.0,<2.0.0", "cupy-cuda13x[ctk]>=14.0.1,!=14.1.0", "libnvforest==26.10.*,>=0.0.0a0", "numpy>=2.0,<3.0a0", - "pylibraft==26.10.*,>=0.0.0a0", "scikit-learn>=1.5", "treelite>=4.7.0,<5.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. @@ -95,10 +95,7 @@ requires = [ "cuda-bindings>=13.0.1,<14.0", "cython>=3.0.0", "libnvforest==26.10.*,>=0.0.0a0", - "libraft==26.10.*,>=0.0.0a0", - "librmm==26.10.*,>=0.0.0a0", "ninja", - "pylibraft==26.10.*,>=0.0.0a0", "treelite>=4.7.0,<5.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. diff --git a/python/nvforest/tests/test_nvforest.py b/python/nvforest/tests/test_nvforest.py index a7282ae..69da752 100644 --- a/python/nvforest/tests/test_nvforest.py +++ b/python/nvforest/tests/test_nvforest.py @@ -10,6 +10,7 @@ import pandas as pd import pytest import treelite +from cuda.core import Device, Stream # Import XGBoost before scikit-learn to work around a libgomp bug # See https://github.com/dmlc/xgboost/issues/7110 @@ -889,6 +890,36 @@ def test_device_selection(device_id, model_kind, tmp_path): assert cp.cuda.runtime.getDevice() == current_device +def test_cuda_core_stream(): + previous_device = Device() + try: + device = Device(0) + device.set_current() + stream = device.create_stream() + + X, y = _simulate_data( + 100, 8, n_informative=2, random_state=0, classification=True + ) + skl_model = RandomForestClassifier( + max_depth=3, random_state=0, n_estimators=5 + ).fit(X, y) + default_fm = nvforest.load_from_sklearn( + skl_model, device="gpu", device_id=0 + ) + fm = nvforest.load_from_sklearn( + skl_model, device="gpu", device_id=0, stream=stream + ) + + assert isinstance(default_fm.forest.stream, Stream) + assert isinstance(stream, Stream) + assert fm.forest.stream is stream + np.testing.assert_array_equal( + cp.asnumpy(fm.predict(X)), skl_model.predict(X) + ) + finally: + previous_device.set_current() + + def test_wide_data(): n_rows = 50 n_features = 100000 From 1f777fc132f21b124c875d287d07039fc3928aa7 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 00:43:51 -0700 Subject: [PATCH 04/17] Update docs --- README.md | 20 +++++++++++ cpp/include/nvforest/Implementation.md | 9 +++-- cpp/include/nvforest/README.md | 20 +++++------ docs/source/build.rst | 16 ++++----- docs/source/conf.py | 4 +++ docs/source/cpp_api.rst | 2 ++ docs/source/getting_started.rst | 49 +++++++++++++++++++++----- docs/source/python_api.rst | 11 ++++++ 8 files changed, 100 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index 1183012..6078602 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,26 @@ fm = nvforest.load_model("/path/to/xgboost_model.ubj", device="gpu", pred = fm.predict(X) ``` +To run GPU work on a caller-managed CUDA stream, pass a `cuda.core.Stream` +created on the model's device. If omitted, nvForest creates and retains a stream +for the model: + +```python +from cuda.core import Device + +device = Device(0) +device.set_current() +stream = device.create_stream() + +fm = nvforest.load_model( + "/path/to/xgboost_model.ubj", + device="gpu", + device_id=0, + stream=stream, +) +pred = fm.predict(X) +``` + Load a scikit-learn random forest model and get class probabilities: ```python diff --git a/cpp/include/nvforest/Implementation.md b/cpp/include/nvforest/Implementation.md index b0bc595..dfa1447 100644 --- a/cpp/include/nvforest/Implementation.md +++ b/cpp/include/nvforest/Implementation.md @@ -147,13 +147,16 @@ a standard benchmark) on the CPU. With some motivation for the general approach to CPU-GPU interoperability, we now offer an overview of the layout of the codebase to help guide future -improvements. Because `nvforest::detail` utilities are going to be moved to RAFT or other -general-purpose libraries, we will not review anything within the `nvforest::detail` -directory here. +improvements. The `nvforest::detail` namespace is private implementation infrastructure and +may change without notice, so this overview focuses on its responsibilities +rather than treating individual utilities as stable APIs. ### Public Headers * `constants.hpp`: Contains constant values that may be useful in working with nvForest in other C++ applications +* `cuda_stream.hpp`: Defines the public `nvforest::cuda_stream` abstraction. + It aliases `cudaStream_t` in GPU builds and is a CUDA-free placeholder in + CPU-only builds. * `decision_forest.hpp`: Provides `decision_forest`, a template which provides concrete implementations of a decision forest. Because different types may be optimal for different sizes of models or models with different features, diff --git a/cpp/include/nvforest/README.md b/cpp/include/nvforest/README.md index 711ab33..588aa84 100644 --- a/cpp/include/nvforest/README.md +++ b/cpp/include/nvforest/README.md @@ -99,25 +99,18 @@ cudaMalloc((void**)&output, num_rows * num_outputs * sizeof(float)); // Assuming that input is a float* pointing to data already located on-device -auto handle = nvforest::handle_t{}; - nvforest_model.predict( - handle, output, input, num_rows, nvforest::device_type::gpu, // out_mem_type nvforest::device_type::gpu, // in_mem_type + stream, + nvforest::infer_kind::default_kind, 4 // chunk_size ); ``` -**handle**: To provide a unified interface on CPU and GPU, we introduce -`nvforest::handle_t` as a wrapper for `raft::handle_t`. This is currently just a -placeholder in CPU-only builds, and using it does not require any CUDA -functionality. For GPU-enabled builds, you can construct a -`nvforest::handle_t` directly from the `raft::handle_t` you wish to use. - **output**: Pointer to pre-allocated buffer where results should be written. If the model has been loaded at single precision, this should be a `float*`. If the model has been loaded at double precision, this should be a @@ -133,6 +126,13 @@ loaded at double precision, this should be a `double*`. **in_mem_type**: Indicates whether input buffer is on device or host. +**stream**: The CUDA stream on which inference and any required data copies +will run. It must be associated with the same GPU device as the model. Reusing +the stream passed during model import preserves the ordering between model +initialization and inference. Synchronize the stream before consuming the +output on the host or releasing resources used by the operation. Keep the stream +alive until the model is destroyed and all queued work has completed. + **chunk_size**: This value has a somewhat different meaning for CPU and GPU, but it generally indicates the number of rows which are evaluated in a single iteration of nvForest's forest evaluation algorithm. On GPU, any power of 2 from 1 to 32 @@ -149,4 +149,4 @@ will be selected based on heuristics. ## Learning More While the above usage summary should be enough to get started using nvForest in another C++ application, you can learn more about the details of this -implementation by reading TODO(wphicks). +implementation in [Implementation.md](Implementation.md). diff --git a/docs/source/build.rst b/docs/source/build.rst index fc4f1e8..d8d0786 100644 --- a/docs/source/build.rst +++ b/docs/source/build.rst @@ -21,14 +21,14 @@ In addition, nvForest requires Cython 3.0 or later. It is possible to build nvForest without GPU support; in this case, the CUDA toolkit is not required. To build nvForest without GPU, set the CMake option ``NVFOREST_ENABLE_GPU=OFF``. -**RAPIDS libraries.** -The nvForest code base is updated in tandem with the rest of RAPIDS. So to build the latest nvForest, you -should use the latest version of RAPIDS as well. (For example, nvForest 26.04 will require 26.04 version of -all RAPIDS packages.) +**Relation to RAPIDS** +The nvForest code base is updated in tandem with the rest of RAPIDS. For example, nvForest 26.08 will be released +together with 26.08 version of other RAPIDS packages. However, it is possible to build and install nvForest +without installing RAPIDS. Only `CCCL `_ is required. **Python dependencies.** -Please see https://docs.rapids.ai/install/ for RAPIDS-wide version support. - +nvForest uses `CUDA Python `_. In addition, nvForest requires +Treelite, cuPy, and scikit-learn. We aim to meet the `SPEC 0 guidelines `_ for minimal supported versions. **For development only.** @@ -292,10 +292,6 @@ nvForest's cmake has the following configurable flags available: - [ON, OFF] - OFF - Whether to cache build artifacts with ccache. - * - NVFOREST_USE_RAFT_STATIC - - [ON, OFF] - - OFF - - Whether to statically link the RAFT library. * - NVFOREST_USE_TREELITE_STATIC - [ON, OFF] - OFF diff --git a/docs/source/conf.py b/docs/source/conf.py index c14d515..6d62048 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -200,6 +200,10 @@ # Example configuration for intersphinx: refer to the Python standard library. intersphinx_mapping = { + "cuda-python": ( + "https://nvidia.github.io/cuda-python/cuda-core/latest/", + None, + ), "numpy": ("https://numpy.org/doc/stable/", None), "python": ("https://docs.python.org/3", None), "sklearn": ("https://scikit-learn.org/stable/", None), diff --git a/docs/source/cpp_api.rst b/docs/source/cpp_api.rst index af145e2..c2d96d9 100644 --- a/docs/source/cpp_api.rst +++ b/docs/source/cpp_api.rst @@ -37,6 +37,8 @@ Enums and constants Type aliases ============ +.. doxygentypedef:: nvforest::cuda_stream + .. doxygentypedef:: nvforest::decision_forest_variant .. doxygentypedef:: nvforest::detail::preset_decision_forest diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 8d48574..127a597 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -128,6 +128,30 @@ You can inspect the type of the model by printing its type: Each model object is associated with a single device. Use the ``device_id`` property to look up which device the model object is located on. +CUDA stream selection +--------------------- + +GPU model loading and inference run on a :external+cuda-python:py:class:`cuda.core.Stream`. +If ``stream`` is omitted, nvForest creates a stream on the selected GPU and keeps it alive +for the lifetime of the model. To supply your own stream, create it on the same device as +``device_id`` and pass it to any model-loading function: + +.. code-block:: python + + from cuda.core import Device + + device = Device(0) + device.set_current() + stream = device.create_stream() + + fm = nvforest.load_from_sklearn( + skl_model, device="gpu", device_id=0, stream=stream + ) + predictions = fm.predict(X) + +nvForest retains a reference to a caller-supplied stream. Python prediction methods synchronize +the model's stream before returning, so their results are ready for immediate use. + Running inference ----------------- @@ -222,29 +246,38 @@ Once the tree model is available as a Treelite object, pass it to the #include #include #include + #include #include + cudaStream_t stream{}; + cudaStreamCreate(&stream); + auto fm = nvforest::import_from_treelite_model( *treelite_model, nvforest::preferred_tree_layout, nvforest::index_type{}, std::nullopt, - nvforest::device_type::gpu); + nvforest::device_type::gpu, + 0, + stream); Now that the tree model is fully imported into nvForest, let's run inference: .. code-block:: cpp - #include - #include - - raft::handle_t raft_handle{}; - nvforest::handle_t handle{raft_handle}; - // Assumption: // * Both output and input are in the GPU memory. // * The input buffer should be of dimension (num_rows, num_features) // * The output buffer should be of dimension (num_rows, fm.num_outputs()) - fm.predict(handle, output, input, num_rows, + fm.predict(output, input, num_rows, nvforest::device_type::gpu, nvforest::device_type::gpu, + stream, nvforest::infer_kind::default_kind); + + // Wait before reading or freeing output. + cudaStreamSynchronize(stream); + +Use a stream associated with the same device as the model. Reusing the stream for import and +inference preserves the ordering between model initialization and prediction without an +additional synchronization. Keep the stream alive until the model is destroyed and all queued +work has completed. diff --git a/docs/source/python_api.rst b/docs/source/python_api.rst index 68cbf0a..3aeea84 100644 --- a/docs/source/python_api.rst +++ b/docs/source/python_api.rst @@ -2,6 +2,17 @@ Python API ########## +CUDA streams +============ + +GPU models use :external+cuda-python:py:class:`cuda.core.Stream` for model loading and +inference. Pass a stream with the ``stream`` parameter on any loading function, or omit it +to let nvForest create and retain a stream on the selected GPU. A supplied stream must be +associated with the same device as ``device_id``. + +Model loading functions +======================= + .. role:: py(code) :language: python :class: highlight From 1bc35da03051eb1a7db31131e6a98c2e5bd05f5a Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 00:54:25 -0700 Subject: [PATCH 05/17] coderrabbitai: Add stream arg to owning buffers --- cpp/include/nvforest/forest_model.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index f882d4c..b60edf3 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -153,14 +153,14 @@ struct forest_model { auto row_count = input.size() / num_features(); auto buffer_in = buffer{}; if (input.memory_type() != memory_type()) { - buffer_in = buffer{row_count * num_features(), memory_type()}; + buffer_in = buffer{row_count * num_features(), memory_type(), stream}; copy(buffer_in, input, 0, 0, buffer_in.size(), stream); } else { buffer_in = buffer{input.data(), row_count * num_features(), memory_type()}; } auto buffer_out = buffer{}; if (output.memory_type() != memory_type()) { - buffer_out = buffer{row_count * num_outputs(), memory_type()}; + buffer_out = buffer{row_count * num_outputs(), memory_type(), stream}; } else { buffer_out = buffer{output.data(), row_count * num_outputs(), memory_type()}; } From aec2bdca8d77fc3443c68cf54eec36a0a96ad6f8 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 00:59:04 -0700 Subject: [PATCH 06/17] Fix buffer construction --- cpp/include/nvforest/forest_model.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index b60edf3..be305f4 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -153,14 +153,16 @@ struct forest_model { auto row_count = input.size() / num_features(); auto buffer_in = buffer{}; if (input.memory_type() != memory_type()) { - buffer_in = buffer{row_count * num_features(), memory_type(), stream}; + buffer_in = + buffer{row_count * num_features(), memory_type(), device_index(), stream}; copy(buffer_in, input, 0, 0, buffer_in.size(), stream); } else { buffer_in = buffer{input.data(), row_count * num_features(), memory_type()}; } auto buffer_out = buffer{}; if (output.memory_type() != memory_type()) { - buffer_out = buffer{row_count * num_outputs(), memory_type(), stream}; + buffer_out = + buffer{row_count * num_outputs(), memory_type(), device_index(), stream}; } else { buffer_out = buffer{output.data(), row_count * num_outputs(), memory_type()}; } From 1fc19969332ba7818de7b3fb86f885818d2e958d Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 16:04:00 -0700 Subject: [PATCH 07/17] Implement stream protocol --- README.md | 2 +- python/nvforest/nvforest/__init__.py | 2 ++ python/nvforest/nvforest/_factory.py | 16 +++++++------- python/nvforest/nvforest/_forest_inference.py | 21 +++++++++---------- python/nvforest/nvforest/_typing.py | 11 ++++++++-- .../nvforest/detail/forest_inference.pyx | 19 ++++++++++------- 6 files changed, 41 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 6078602..ab32425 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ fm = nvforest.load_model( device_id=0, stream=stream, ) -pred = fm.predict(X) +pred = fm.predict(X) # Note: automatically syncs the stream ``` Load a scikit-learn random forest model and get class probabilities: diff --git a/python/nvforest/nvforest/__init__.py b/python/nvforest/nvforest/__init__.py index 3fe9542..e5a0130 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -24,6 +24,7 @@ GPUForestInferenceClassifier, GPUForestInferenceRegressor, ) +from nvforest._typing import StreamLike from nvforest._version import __git_commit__, __version__ __all__ = [ @@ -31,6 +32,7 @@ "CPUForestInferenceRegressor", "GPUForestInferenceClassifier", "GPUForestInferenceRegressor", + "StreamLike", "load_model", "load_from_sklearn", "load_from_treelite_model", diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index 37da9a7..b999917 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -7,7 +7,6 @@ from typing import Any, Optional, Union import treelite -from cuda.core import Stream from nvforest._base import ForestInference from nvforest._forest_inference import ( @@ -18,6 +17,7 @@ infer_device, infer_is_classifier, ) +from nvforest._typing import StreamLike def get_forest_inference_class(device, is_classifier) -> type: @@ -39,7 +39,7 @@ def make_forest_inference_object( treelite_model: treelite.Model, device: str, device_id: Optional[int], - stream: Optional[Stream], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -72,7 +72,7 @@ def load_model( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, ) -> ForestInference: """Load a model into nvForest from a serialized model file. @@ -113,7 +113,7 @@ def load_model( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - stream : cuda.core.Stream or None, default=None + stream : StreamLike or None, default=None For GPU execution, the CUDA stream to use during model loading and inference. If not given, a new stream will be created. For CPU execution, this value is ignored. @@ -174,7 +174,7 @@ def load_from_sklearn( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, ) -> ForestInference: """Load a Scikit-Learn forest model to nvForest @@ -208,7 +208,7 @@ def load_from_sklearn( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - stream : cuda.core.Stream or None, default=None + stream : StreamLike or None, default=None For GPU execution, the CUDA stream to use during model loading and inference. If not given, a new stream will be created. For CPU execution, this value is ignored. @@ -236,7 +236,7 @@ def load_from_treelite_model( align_bytes: Optional[int] = None, precision: Optional[str] = None, device_id: Optional[int] = None, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, ) -> ForestInference: """Load a Treelite forest model to nvForest @@ -270,7 +270,7 @@ def load_from_treelite_model( device_id : int or None, default=None For GPU execution, the device on which to load and execute this model. For CPU execution, this value is currently ignored. - stream : cuda.core.Stream or None, default=None + stream : StreamLike or None, default=None For GPU execution, the CUDA stream to use during model loading and inference. If not given, a new stream will be created. For CPU execution, this value is ignored. diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index e312776..2d67e87 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -16,10 +16,9 @@ import numpy as np import treelite from cuda.bindings import runtime -from cuda.core import Stream from nvforest._base import ForestInferenceClassifier, ForestInferenceRegressor -from nvforest._typing import DataType +from nvforest._typing import DataType, StreamLike from nvforest.detail.forest_inference import ForestInferenceImpl @@ -123,7 +122,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[Stream], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -260,7 +259,7 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -284,7 +283,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[Stream], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -397,7 +396,7 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -421,7 +420,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[Stream], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -526,7 +525,7 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -551,7 +550,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[Stream], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -665,7 +664,7 @@ def __init__( self, *, treelite_model: treelite.Model, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -690,7 +689,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - stream: Optional[Stream], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], diff --git a/python/nvforest/nvforest/_typing.py b/python/nvforest/nvforest/_typing.py index 548a05b..5f897fb 100644 --- a/python/nvforest/nvforest/_typing.py +++ b/python/nvforest/nvforest/_typing.py @@ -1,9 +1,9 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Protocol, Union, runtime_checkable import numpy as np @@ -11,3 +11,10 @@ import cupy DataType = Union[np.ndarray, "cupy.ndarray"] + + +@runtime_checkable +class StreamLike(Protocol): + """Duck typing for all stream-like objects""" + + def __cuda_stream__(self) -> tuple[int, int]: ... diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index ced0918..61aa29a 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -7,9 +7,9 @@ from typing import Optional, Union import numpy as np import treelite -from cuda.core import Device, Stream +from cuda.core import Device -from nvforest._typing import DataType +from nvforest._typing import DataType, StreamLike from nvforest.detail.treelite import safe_treelite_call from libc.stdint cimport uint32_t, uintptr_t @@ -78,11 +78,14 @@ cdef class ForestInference_impl(): device: str = "cpu", device_id: Optional[int] = None, ): + # Assumption: The caller needs to pass in correct (device, device_id) pair + # This function will not contain any logic for auto-detecting device. cdef uintptr_t stream_ptr = 0 if stream is not None: - if not isinstance(stream, Stream): - raise TypeError("stream must be a cuda.core.Stream or None") - stream_ptr = int(stream.handle) + if not isinstance(stream, StreamLike): + raise TypeError("stream must be a stream-like object or None") + stream_tuple: tuple[int, int] = stream.__cuda_stream__() + stream_ptr = stream_tuple[1] self.stream = stream self.stream_handle = stream_ptr @@ -272,7 +275,7 @@ class ForestInferenceImpl: treelite_model: treelite.Model, device: str, device_id: int, - stream: Optional[Stream] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -280,8 +283,8 @@ class ForestInferenceImpl: ): # Assumption: The caller needs to pass in correct (device, device_id) pair # This function will not contain any logic for auto-detecting device. - if stream is not None and not isinstance(stream, Stream): - raise TypeError("stream must be a cuda.core.Stream or None") + if stream is not None and not isinstance(stream, StreamLike): + raise TypeError("stream must be a stream-like object or None") if device == "gpu" and stream is None: previous_device = Device() try: From 6f3259885e02054a202bc68532e52723d11c7e0c Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 16:10:43 -0700 Subject: [PATCH 08/17] Update pytest --- python/nvforest/tests/test_nvforest.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/python/nvforest/tests/test_nvforest.py b/python/nvforest/tests/test_nvforest.py index 69da752..4bb4fbd 100644 --- a/python/nvforest/tests/test_nvforest.py +++ b/python/nvforest/tests/test_nvforest.py @@ -798,22 +798,13 @@ def test_missing_categorical(category_list): input = np.array([[np.nan]]) gtil_preds = treelite.gtil.predict(model, input) fm = nvforest.load_from_treelite_model(model) - fil_preds = _get_numpy_array(fm.predict(input)) - np.testing.assert_equal(fil_preds.flatten(), gtil_preds.flatten()) + nvforest_preds = _get_numpy_array(fm.predict(input)) + np.testing.assert_equal(nvforest_preds.flatten(), gtil_preds.flatten()) @pytest.mark.parametrize("device_id", [None, 0, 1, 2]) @pytest.mark.parametrize("model_kind", ["sklearn", "xgboost"]) def test_device_selection(device_id, model_kind, tmp_path): - import sklearn - from packaging.version import Version - - # TODO(hcho3): Remove this once Rapids adopts XGBoost 3.1.3 - if model_kind == "xgboost" and Version(sklearn.__version__) >= Version( - "1.8.0.dev0" - ): - pytest.skip("xgboost is incompatible with sklearn >= 1.8.0.dev0") - current_device = cp.cuda.runtime.getDevice() if device_id is not None and device_id >= cp.cuda.runtime.getDeviceCount(): From 0aab12a532dff0c61d5522653029bba3d101e65b Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 17:02:09 -0700 Subject: [PATCH 09/17] Move buffers instead of copying, to preserve stream assignment --- cpp/include/nvforest/decision_forest.hpp | 19 ++++++++++--------- cpp/include/nvforest/forest_model.hpp | 3 ++- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/cpp/include/nvforest/decision_forest.hpp b/cpp/include/nvforest/decision_forest.hpp index f982338..61747e5 100644 --- a/cpp/include/nvforest/decision_forest.hpp +++ b/cpp/include/nvforest/decision_forest.hpp @@ -26,6 +26,7 @@ #include #include #include +#include #include namespace nvforest { @@ -161,12 +162,12 @@ struct decision_forest { element_op elem_postproc = element_op::disable, io_type average_factor = io_type{1}, io_type postproc_constant = io_type{1}) - : nodes_{nodes}, - root_node_indexes_{root_node_indexes}, - node_id_mapping_{node_id_mapping}, - bias_{bias}, - vector_output_{vector_output}, - categorical_storage_{categorical_storage}, + : nodes_{std::move(nodes)}, + root_node_indexes_{std::move(root_node_indexes)}, + node_id_mapping_{std::move(node_id_mapping)}, + bias_{std::move(bias)}, + vector_output_{std::move(vector_output)}, + categorical_storage_{std::move(categorical_storage)}, num_features_{num_features}, num_outputs_{num_outputs}, leaf_size_{leaf_size}, @@ -176,15 +177,15 @@ struct decision_forest { average_factor_{average_factor}, postproc_constant_{postproc_constant} { - if (nodes.memory_type() != root_node_indexes.memory_type()) { + if (nodes_.memory_type() != root_node_indexes_.memory_type()) { throw detail::mem_type_mismatch( "Nodes and indexes of forest must both be stored on either host or device"); } - if (nodes.device_index() != root_node_indexes.device_index()) { + if (nodes_.device_index() != root_node_indexes_.device_index()) { throw detail::mem_type_mismatch( "Nodes and indexes of forest must both be stored on same device"); } - detail::initialize_device(nodes.device()); + detail::initialize_device(nodes_.device()); } /** The number of features per row expected by the model */ diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index be305f4..b504d6d 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -19,6 +19,7 @@ #include #include +#include #include namespace nvforest { @@ -32,7 +33,7 @@ namespace nvforest { struct forest_model { /** Wrap a decision_forest in a full forest_model object */ forest_model(decision_forest_variant&& forest = decision_forest_variant{}) - : decision_forest_{forest} + : decision_forest_{std::move(forest)} { } From b1c70c2800060ccb7b63431efe17f2a8ca999c70 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 17:05:39 -0700 Subject: [PATCH 10/17] Update docs for streams --- README.md | 10 +++++++--- docs/source/build.rst | 7 +------ docs/source/getting_started.rst | 11 +++++++---- docs/source/python_api.rst | 8 -------- 4 files changed, 15 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index ab32425..aa10290 100644 --- a/README.md +++ b/README.md @@ -17,9 +17,13 @@ fm = nvforest.load_model("/path/to/xgboost_model.ubj", device="gpu", pred = fm.predict(X) ``` -To run GPU work on a caller-managed CUDA stream, pass a `cuda.core.Stream` -created on the model's device. If omitted, nvForest creates and retains a stream -for the model: +By default, nvForest creates a new +[CUDA stream](https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/asynchronous-execution.html#cuda-streams) +for the model and retains it. To use a custom stream, pass in a :external+cuda-python:py:class:`cuda.core.Stream` or +other stream-like objects (*): + +(*) A stream-like object is an object that exposes a method named `__cuda_stream__`. See [Stream Protocol]( +https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol) for more details. ```python from cuda.core import Device diff --git a/docs/source/build.rst b/docs/source/build.rst index d8d0786..c854fe8 100644 --- a/docs/source/build.rst +++ b/docs/source/build.rst @@ -24,12 +24,7 @@ In addition, nvForest requires Cython 3.0 or later. **Relation to RAPIDS** The nvForest code base is updated in tandem with the rest of RAPIDS. For example, nvForest 26.08 will be released together with 26.08 version of other RAPIDS packages. However, it is possible to build and install nvForest -without installing RAPIDS. Only `CCCL `_ is required. - -**Python dependencies.** -nvForest uses `CUDA Python `_. In addition, nvForest requires -Treelite, cuPy, and scikit-learn. -We aim to meet the `SPEC 0 guidelines `_ for minimal supported versions. +without installing RAPIDS. **For development only.** diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 127a597..e2b45db 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -131,10 +131,10 @@ You can inspect the type of the model by printing its type: CUDA stream selection --------------------- -GPU model loading and inference run on a :external+cuda-python:py:class:`cuda.core.Stream`. -If ``stream`` is omitted, nvForest creates a stream on the selected GPU and keeps it alive -for the lifetime of the model. To supply your own stream, create it on the same device as -``device_id`` and pass it to any model-loading function: +By default, nvForest creates a new +`CUDA stream `_ +for a GPU model. To use a custom stream, pass in a :external+cuda-python:py:class:`cuda.core.Stream` (or +any stream-like object [#]_) via the ``stream`` parameter. A supplied stream must be associated with the same device as ``device_id``. .. code-block:: python @@ -152,6 +152,9 @@ for the lifetime of the model. To supply your own stream, create it on the same nvForest retains a reference to a caller-supplied stream. Python prediction methods synchronize the model's stream before returning, so their results are ready for immediate use. +.. [#] A stream-like object is an object that exposes a method named `__cuda_stream__`. + See `Stream Protocol `_ for more details. + Running inference ----------------- diff --git a/docs/source/python_api.rst b/docs/source/python_api.rst index 3aeea84..8abcbff 100644 --- a/docs/source/python_api.rst +++ b/docs/source/python_api.rst @@ -2,14 +2,6 @@ Python API ########## -CUDA streams -============ - -GPU models use :external+cuda-python:py:class:`cuda.core.Stream` for model loading and -inference. Pass a stream with the ``stream`` parameter on any loading function, or omit it -to let nvForest create and retain a stream on the selected GPU. A supplied stream must be -associated with the same device as ``device_id``. - Model loading functions ======================= From bfb7327d6f3170a545bc30294d4161f48967e6a4 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 17:21:23 -0700 Subject: [PATCH 11/17] Show stream delete in the doc --- docs/source/getting_started.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index e2b45db..245163d 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -280,6 +280,14 @@ Now that the tree model is fully imported into nvForest, let's run inference: // Wait before reading or freeing output. cudaStreamSynchronize(stream); + // Use output + // ... + + // Clean up + fm.reset(); + cudaStreamDestroy(stream); + // ... also free input, output buffers + Use a stream associated with the same device as the model. Reusing the stream for import and inference preserves the ordering between model initialization and prediction without an additional synchronization. Keep the stream alive until the model is destroyed and all queued From 729d37ebb48ee904fdc6a02bf139fc1649cbbf74 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 12 Aug 2026 17:25:26 -0700 Subject: [PATCH 12/17] cimport cudaStream_t from cuda.bindings --- python/nvforest/nvforest/detail/cuda_stream.pxd | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/nvforest/nvforest/detail/cuda_stream.pxd b/python/nvforest/nvforest/detail/cuda_stream.pxd index 807030b..0fea6ca 100644 --- a/python/nvforest/nvforest/detail/cuda_stream.pxd +++ b/python/nvforest/nvforest/detail/cuda_stream.pxd @@ -2,10 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -cdef extern from "cuda_runtime_api.h" nogil: - ctypedef struct CUstream_st: - pass - ctypedef CUstream_st* cudaStream_t +from cuda.bindings.cyruntime cimport cudaStream_t cdef extern from "nvforest/cuda_stream.hpp" namespace "nvforest" nogil: From 29f4bac78aea2ea5447dce9180e817b28d2b7e55 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 00:21:14 -0700 Subject: [PATCH 13/17] Link with CCCL --- cpp/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index f8d2101..c8d4113 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -309,7 +309,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE) endif() if(NVFOREST_ENABLE_GPU) - list(APPEND _nvforest_cpp_public_libs CUDA::cudart_static) + list(APPEND _nvforest_cpp_public_libs CUDA::cudart_static CCCL::CCCL) endif() # These are always private: From 2cd4f11a153ff1e363d894bfdd6c1c75a0a303ae Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 19 Oct 2026 19:44:37 -0700 Subject: [PATCH 14/17] Remove unneeded changes --- ci/build_wheel_libnvforest.sh | 1 + ci/build_wheel_nvforest.sh | 1 + ci/release/update-version.sh | 3 + conda/recipes/nvforest/recipe.yaml | 2 + cpp/cmake/thirdparty/get_raft.cmake | 78 ++++++++++++++++++++++++++ docs/source/conf.py | 8 +-- python/libnvforest/libnvforest/load.py | 2 + python/nvforest/nvforest/_handle.py | 9 +++ 8 files changed, 100 insertions(+), 4 deletions(-) create mode 100644 cpp/cmake/thirdparty/get_raft.cmake create mode 100644 python/nvforest/nvforest/_handle.py diff --git a/ci/build_wheel_libnvforest.sh b/ci/build_wheel_libnvforest.sh index d03ca41..2ffc303 100755 --- a/ci/build_wheel_libnvforest.sh +++ b/ci/build_wheel_libnvforest.sh @@ -29,6 +29,7 @@ rapids-pip-retry install \ export PIP_NO_BUILD_ISOLATION=0 EXCLUDE_ARGS=( + --exclude "libraft.so" --exclude "libcublas.so.*" --exclude "libcublasLt.so.*" --exclude "libcufft.so.*" diff --git a/ci/build_wheel_nvforest.sh b/ci/build_wheel_nvforest.sh index e86fe3d..3d785d0 100755 --- a/ci/build_wheel_nvforest.sh +++ b/ci/build_wheel_nvforest.sh @@ -21,6 +21,7 @@ echo "libnvforest-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBNVFOREST_WHEELH EXCLUDE_ARGS=( --exclude "libnvforest.so" + --exclude "libraft.so" --exclude "libcublas.so.*" --exclude "libcublasLt.so.*" --exclude "libcufft.so.*" diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index 42e57b9..b18d895 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -109,6 +109,9 @@ DEPENDENCIES=( nvforest libnvforest libnvforest-tests + libraft + libraft-headers + pylibraft ) for DEP in "${DEPENDENCIES[@]}"; do for FILE in dependencies.yaml conda/environments/*.yaml; do diff --git a/conda/recipes/nvforest/recipe.yaml b/conda/recipes/nvforest/recipe.yaml index fef271f..d4896bf 100644 --- a/conda/recipes/nvforest/recipe.yaml +++ b/conda/recipes/nvforest/recipe.yaml @@ -81,6 +81,7 @@ requirements: - cython >=3.0.0 - libnvforest =${{ version }} - pip + - pylibraft =${{ minor_version }} - cuda-core >=1.0.0,<2.0.0 - python =${{ py_abi_min }} - python-abi3 ${{ py_abi_min }}.* @@ -98,6 +99,7 @@ requirements: - librmm =${{ minor_version }} - numpy >=2.0,<3.0a0 - scikit-learn >=1.4 + - pylibraft =${{ minor_version }} - cuda-core >=1.0.0,<2.0.0 - python - treelite ${{ treelite_version }} diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake new file mode 100644 index 0000000..1265804 --- /dev/null +++ b/cpp/cmake/thirdparty/get_raft.cmake @@ -0,0 +1,78 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +set(NVFOREST_MIN_VERSION_raft "${NVForest_VERSION_MAJOR}.${NVForest_VERSION_MINOR}.00") + +function(find_and_configure_raft) + set(oneValueArgs + VERSION + FORK + PINNED_TAG + EXCLUDE_FROM_ALL + USE_RAFT_STATIC + COMPILE_LIBRARY + CLONE_ON_PIN + NVTX) + cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "${rapids-cmake-checkout-tag}") + message(STATUS "NVFOREST: RAFT pinned tag found: ${PKG_PINNED_TAG}. Cloning raft locally.") + set(CPM_DOWNLOAD_raft ON) + elseif(PKG_USE_RAFT_STATIC AND (NOT CPM_raft_SOURCE)) + message(STATUS "NVFOREST: Cloning raft locally to build static libraries.") + set(CPM_DOWNLOAD_raft ON) + endif() + + # We need RAFT::distributed for MG tests + if(BUILD_NVFOREST_MG_TESTS) + string(APPEND RAFT_COMPONENTS " distributed") + endif() + + # We need to set this each time so that on subsequent calls to cmake the raft-config.cmake re-evaluates the RAFT_NVTX + # value + set(RAFT_NVTX ${PKG_NVTX}) + + message(VERBOSE "NVFOREST: raft FIND_PACKAGE_ARGUMENTS COMPONENTS ${RAFT_COMPONENTS}") + + rapids_cpm_find( + raft ${PKG_VERSION} + GLOBAL_TARGETS raft::raft + BUILD_EXPORT_SET nvforest-exports + INSTALL_EXPORT_SET nvforest-exports COMPONENTS ${RAFT_COMPONENTS} + CPM_ARGS + GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git + GIT_TAG ${PKG_PINNED_TAG} SOURCE_SUBDIR cpp + EXCLUDE_FROM_ALL ${PKG_EXCLUDE_FROM_ALL} + OPTIONS "BUILD_TESTS OFF" "BUILD_PRIMS_BENCH OFF" "BUILD_CAGRA_HNSWLIB OFF" "RAFT_COMPILE_LIBRARY OFF") + + if(raft_ADDED) + message(VERBOSE "NVFOREST: Using RAFT located in ${raft_SOURCE_DIR}") + else() + message(VERBOSE "NVFOREST: Using RAFT located in ${raft_DIR}") + endif() + +endfunction() + +# Change pinned tag here to test a commit in CI To use a different RAFT locally, set the CMake variable +# CPM_raft_SOURCE=/path/to/local/raft +find_and_configure_raft( + VERSION + ${NVFOREST_MIN_VERSION_raft} + FORK + rapidsai + PINNED_TAG + ${rapids-cmake-checkout-tag} + EXCLUDE_FROM_ALL + ${NVFOREST_EXCLUDE_RAFT_FROM_ALL} + # When PINNED_TAG above doesn't match nvforest, force local raft clone in build directory even if it's already + # installed. + CLONE_ON_PIN + ${NVFOREST_RAFT_CLONE_ON_PIN} + USE_RAFT_STATIC + ${NVFOREST_USE_RAFT_STATIC} + NVTX + ${NVTX}) diff --git a/docs/source/conf.py b/docs/source/conf.py index d32f85c..f0aedd0 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -201,14 +201,14 @@ # Example configuration for intersphinx: refer to the Python standard library. intersphinx_mapping = { - "cuda-python": ( - "https://nvidia.github.io/cuda-python/cuda-core/latest/", - None, - ), "numpy": ("https://numpy.org/doc/stable/", None), "python": ("https://docs.python.org/3", None), "sklearn": ("https://scikit-learn.org/stable/", None), "treelite": ("https://treelite.readthedocs.io/en/latest/", None), + "cuda-python": ( + "https://nvidia.github.io/cuda-python/cuda-core/", + None, + ), } # Config numpydoc diff --git a/python/libnvforest/libnvforest/load.py b/python/libnvforest/libnvforest/load.py index 88fede2..95b837a 100644 --- a/python/libnvforest/libnvforest/load.py +++ b/python/libnvforest/libnvforest/load.py @@ -37,9 +37,11 @@ def load_library(): """Dynamically load libnvforest.so and its dependencies""" try: # These libraries must all be loaded before libnvforest + import libraft import rapids_logger rapids_logger.load_library() + libraft.load_library() except ModuleNotFoundError: # These runtime dependencies might be satisfied by conda packages (which do not # have any Python modules) instead of wheels. In that situation, assume that diff --git a/python/nvforest/nvforest/_handle.py b/python/nvforest/nvforest/_handle.py new file mode 100644 index 0000000..347b748 --- /dev/null +++ b/python/nvforest/nvforest/_handle.py @@ -0,0 +1,9 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# + +from pylibraft.common.handle import Handle as RaftHandle + +# For now, nvforest.handle.Handle is an alias of pylibraft.common.handle.Handle +Handle = RaftHandle From 527e1589986c5a92f9a67d1549e6665541c57f05 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 19 Oct 2026 20:05:11 -0700 Subject: [PATCH 15/17] Add back librmm to Conda env --- ci/build_wheel_libnvforest.sh | 3 +- ci/build_wheel_nvforest.sh | 1 + ci/release/update-version.sh | 1 + .../all_cuda-129_arch-aarch64.yaml | 1 + .../all_cuda-129_arch-x86_64.yaml | 1 + .../all_cuda-133_arch-aarch64.yaml | 1 + .../all_cuda-133_arch-x86_64.yaml | 1 + .../clang_tidy_cuda-129_arch-x86_64.yaml | 1 + .../clang_tidy_cuda-133_arch-x86_64.yaml | 1 + .../cpp_all_cuda-129_arch-x86_64.yaml | 1 + .../cpp_all_cuda-133_arch-x86_64.yaml | 1 + conda/recipes/libnvforest/recipe.yaml | 6 ++++ dependencies.yaml | 33 +++++++++++++++++++ python/libnvforest/libnvforest/load.py | 2 ++ python/libnvforest/pyproject.toml | 2 ++ python/nvforest/pyproject.toml | 1 + 16 files changed, 56 insertions(+), 1 deletion(-) diff --git a/ci/build_wheel_libnvforest.sh b/ci/build_wheel_libnvforest.sh index 2ffc303..b6b8db8 100755 --- a/ci/build_wheel_libnvforest.sh +++ b/ci/build_wheel_libnvforest.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -38,6 +38,7 @@ EXCLUDE_ARGS=( --exclude "libcusparse.so.*" --exclude "libnvJitLink.so.*" --exclude "librapids_logger.so" + --exclude "librmm.so" ) export SKBUILD_CMAKE_ARGS="-DDISABLE_DEPRECATION_WARNINGS=ON" diff --git a/ci/build_wheel_nvforest.sh b/ci/build_wheel_nvforest.sh index 3d785d0..36c7e8a 100755 --- a/ci/build_wheel_nvforest.sh +++ b/ci/build_wheel_nvforest.sh @@ -30,6 +30,7 @@ EXCLUDE_ARGS=( --exclude "libcusparse.so.*" --exclude "libnvJitLink.so.*" --exclude "librapids_logger.so" + --exclude "librmm.so" --exclude "libtreelite.so" ) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index b18d895..53b82a0 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -111,6 +111,7 @@ DEPENDENCIES=( libnvforest-tests libraft libraft-headers + librmm pylibraft ) for DEP in "${DEPENDENCIES[@]}"; do diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 49753bd..e4216f5 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -23,6 +23,7 @@ dependencies: - gcc_linux-aarch64=14.* - graphviz - libraft==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 7a570c5..6b5d28d 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -23,6 +23,7 @@ dependencies: - gcc_linux-64=14.* - graphviz - libraft==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 3d690c8..ac8b604 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -23,6 +23,7 @@ dependencies: - gcc_linux-aarch64=14.* - graphviz - libraft==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 2ed9130..0491de1 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -23,6 +23,7 @@ dependencies: - gcc_linux-64=14.* - graphviz - libraft==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - ninja - numpy>=2.0,<3.0a0 - numpydoc diff --git a/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml b/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml index 18f8fd2..b57a5eb 100644 --- a/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml +++ b/conda/environments/clang_tidy_cuda-129_arch-x86_64.yaml @@ -16,6 +16,7 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libraft-headers==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - llvm-openmp==15.0.7 - ninja - sysroot_linux-64==2.28 diff --git a/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml b/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml index 1466929..c3e5fa2 100644 --- a/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml +++ b/conda/environments/clang_tidy_cuda-133_arch-x86_64.yaml @@ -16,6 +16,7 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libraft-headers==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - llvm-openmp==15.0.7 - ninja - sysroot_linux-64==2.28 diff --git a/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml b/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml index 205c851..ff41d6c 100644 --- a/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/cpp_all_cuda-129_arch-x86_64.yaml @@ -14,6 +14,7 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libraft-headers==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - ninja - sysroot_linux-64==2.28 name: cpp_all_cuda-129_arch-x86_64 diff --git a/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml b/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml index dce048a..f7e6a7c 100644 --- a/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/cpp_all_cuda-133_arch-x86_64.yaml @@ -14,6 +14,7 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libraft-headers==26.12.*,>=0.0.0a0 +- librmm==26.12.*,>=0.0.0a0 - ninja - sysroot_linux-64==2.28 name: cpp_all_cuda-133_arch-x86_64 diff --git a/conda/recipes/libnvforest/recipe.yaml b/conda/recipes/libnvforest/recipe.yaml index c6948e9..262c7c3 100644 --- a/conda/recipes/libnvforest/recipe.yaml +++ b/conda/recipes/libnvforest/recipe.yaml @@ -77,6 +77,7 @@ outputs: host: - cuda-version =${{ cuda_version }} - libraft-headers =${{ minor_version }} + - librmm =${{ minor_version }} - rapids-logger =0.3 - treelite ${{ treelite_version }} - cuda-cudart-dev @@ -103,17 +104,20 @@ outputs: host: - cuda-version =${{ cuda_version }} - rapids-logger =0.3 + - librmm =${{ minor_version }} - treelite ${{ treelite_version }} - cuda-cudart-dev run: - libgomp >=${{ c_compiler_version }} - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} + - librmm =${{ minor_version }} - rapids-logger =0.3 - treelite ${{ treelite_version }} ignore_run_exports: by_name: - cuda-cudart - cuda-version + - librmm - libcusolver about: homepage: ${{ load_from_file("python/libnvforest/pyproject.toml").project.urls.Homepage }} @@ -141,11 +145,13 @@ outputs: run: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - ${{ pin_subpackage("libnvforest", exact=True) }} + - librmm =${{ minor_version }} - treelite ${{ treelite_version }} ignore_run_exports: by_name: - cuda-cudart - cuda-version + - librmm about: homepage: ${{ load_from_file("python/libnvforest/pyproject.toml").project.urls.Homepage }} license: ${{ load_from_file("python/libnvforest/pyproject.toml").project.license }} diff --git a/dependencies.yaml b/dependencies.yaml index 272a9ec..5bdbfed 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -16,6 +16,7 @@ files: - depends_on_cuda_core - depends_on_cupy - depends_on_libraft + - depends_on_librmm - depends_on_pylibraft - depends_on_rapids_logger - develop @@ -28,6 +29,10 @@ files: devcontainers: output: none includes: + # "devcontainers" includes everything in "all", excluding test_python_xgboost + # The libxgboost package depends on librmm but we do not want to have a + # package depending on librmm in devcontainers since it should be built + # from source. - common_build - cuda - cuda_version @@ -35,6 +40,7 @@ files: - depends_on_cuda_core - depends_on_cupy - depends_on_libraft + - depends_on_librmm - depends_on_pylibraft - depends_on_rapids_logger - py_build_nvforest @@ -52,6 +58,7 @@ files: - cuda - cuda_version - depends_on_libraft_headers + - depends_on_librmm checks: output: none includes: @@ -68,6 +75,7 @@ files: - cuda - cuda_version - depends_on_libraft_headers + - depends_on_librmm docs: output: none includes: @@ -112,6 +120,7 @@ files: - depends_on_cuda_python - depends_on_libnvforest - depends_on_libraft + - depends_on_librmm - depends_on_pylibraft - py_build_nvforest py_run_nvforest: @@ -159,6 +168,7 @@ files: includes: - common_build - depends_on_libraft + - depends_on_librmm - depends_on_rapids_logger py_run_libnvforest: output: pyproject @@ -170,6 +180,7 @@ files: table: project includes: - depends_on_libraft + - depends_on_librmm - depends_on_rapids_logger channels: - rapidsai-nightly @@ -537,6 +548,28 @@ dependencies: - output_types: conda packages: - libraft-headers==26.12.*,>=0.0.0a0 + depends_on_librmm: + common: + - output_types: conda + packages: + - &librmm_unsuffixed librmm==26.12.*,>=0.0.0a0 + specific: + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + cuda_suffixed: "true" + packages: + - librmm-cu12==26.12.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - librmm-cu13==26.12.*,>=0.0.0a0 + - matrix: + cuda_suffixed: "false" + packages: + - *librmm_unsuffixed depends_on_pylibraft: common: - output_types: conda diff --git a/python/libnvforest/libnvforest/load.py b/python/libnvforest/libnvforest/load.py index 95b837a..462d1fc 100644 --- a/python/libnvforest/libnvforest/load.py +++ b/python/libnvforest/libnvforest/load.py @@ -38,9 +38,11 @@ def load_library(): try: # These libraries must all be loaded before libnvforest import libraft + import librmm import rapids_logger rapids_logger.load_library() + librmm.load_library() libraft.load_library() except ModuleNotFoundError: # These runtime dependencies might be satisfied by conda packages (which do not diff --git a/python/libnvforest/pyproject.toml b/python/libnvforest/pyproject.toml index 7257233..6c888d3 100644 --- a/python/libnvforest/pyproject.toml +++ b/python/libnvforest/pyproject.toml @@ -26,6 +26,7 @@ classifiers = [ ] dependencies = [ "libraft==26.12.*,>=0.0.0a0", + "librmm==26.12.*,>=0.0.0a0", "rapids-logger==0.3.*", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. @@ -64,6 +65,7 @@ matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true" requires = [ "cmake>=4.0", "libraft==26.12.*,>=0.0.0a0", + "librmm==26.12.*,>=0.0.0a0", "ninja", "rapids-logger==0.3.*", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. diff --git a/python/nvforest/pyproject.toml b/python/nvforest/pyproject.toml index 58ae573..24a2514 100644 --- a/python/nvforest/pyproject.toml +++ b/python/nvforest/pyproject.toml @@ -97,6 +97,7 @@ requires = [ "cython>=3.0.0", "libnvforest==26.12.*,>=0.0.0a0", "libraft==26.12.*,>=0.0.0a0", + "librmm==26.12.*,>=0.0.0a0", "ninja", "pylibraft==26.12.*,>=0.0.0a0", "treelite>=4.7.2,<5.0", From 1f6e2264c066496600c300b82904b2ad5682cb23 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 20 Oct 2026 11:57:39 -0700 Subject: [PATCH 16/17] Remove changes in update-version.sh --- ci/release/update-version.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index 53b82a0..bd552e7 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -113,6 +113,7 @@ DEPENDENCIES=( libraft-headers librmm pylibraft + rmm ) for DEP in "${DEPENDENCIES[@]}"; do for FILE in dependencies.yaml conda/environments/*.yaml; do From ce27061c20769651c0b7f8c520b94ed771ae80a9 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Fri, 30 Oct 2026 14:09:51 -0700 Subject: [PATCH 17/17] cuda-core is a runtime dependency --- conda/recipes/nvforest/recipe.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/conda/recipes/nvforest/recipe.yaml b/conda/recipes/nvforest/recipe.yaml index fedf4a3..a935b1f 100644 --- a/conda/recipes/nvforest/recipe.yaml +++ b/conda/recipes/nvforest/recipe.yaml @@ -82,7 +82,6 @@ requirements: - libnvforest =${{ version }} - pip - pylibraft =${{ minor_version }} - - cuda-core >=1.0.0,<2.0.0 - python =${{ py_abi_min }} - python-abi3 ${{ py_abi_min }}.* - rapids-build-backend >=0.4.0,<0.5.0.dev0