From 5de7aebd112433316c97973fbc3c63a087a28142 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 00:43:30 -0700 Subject: [PATCH 01/13] Remove RMM/RAFT deps Part 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 + 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 | 130 ++++-------------- cpp/include/nvforest/handle.hpp | 43 ------ 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 +-- dependencies.yaml | 7 + python/nvforest/nvforest/__init__.py | 6 +- python/nvforest/nvforest/_factory.py | 44 +++--- python/nvforest/nvforest/_forest_inference.py | 43 +++--- python/nvforest/nvforest/_handle.py | 9 -- python/nvforest/nvforest/_typing.py | 11 +- .../nvforest/nvforest/detail/cuda_stream.pxd | 6 +- .../nvforest/detail/forest_inference.pyx | 59 +++++--- python/nvforest/nvforest/detail/handle.pxd | 17 --- 25 files changed, 237 insertions(+), 382 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 delete mode 100644 python/nvforest/nvforest/_handle.py delete mode 100644 python/nvforest/nvforest/detail/handle.pxd diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 6dd75f5..056840f 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 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index d9e43d6..915d79e 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 diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index fdbe67e..16caada 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 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index eb3d2c2..b43abc1 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 diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 885fa54..c8d4113 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 CCCL::CCCL) 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..b504d6d 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 @@ -21,6 +19,7 @@ #include #include +#include #include namespace nvforest { @@ -34,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)} { } @@ -146,105 +145,32 @@ 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(), 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(), device_index(), stream}; + } 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 +183,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 +190,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 +209,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 +230,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 fd1f3bb..0000000 --- a/cpp/include/nvforest/handle.hpp +++ /dev/null @@ -1,43 +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 - { - 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_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(); - } - - 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); diff --git a/dependencies.yaml b/dependencies.yaml index ddcbbe2..d85a743 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,6 +13,7 @@ files: - cuda - cuda_version - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libraft - depends_on_librmm @@ -37,6 +38,7 @@ files: - cuda - cuda_version - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libraft - depends_on_librmm @@ -455,6 +457,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] diff --git a/python/nvforest/nvforest/__init__.py b/python/nvforest/nvforest/__init__.py index 4fc9b15..e5a0130 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,7 @@ GPUForestInferenceClassifier, GPUForestInferenceRegressor, ) -from nvforest._handle import Handle +from nvforest._typing import StreamLike from nvforest._version import __git_commit__, __version__ __all__ = [ @@ -32,7 +32,7 @@ "CPUForestInferenceRegressor", "GPUForestInferenceClassifier", "GPUForestInferenceRegressor", - "Handle", + "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 cc0700a..b999917 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 # @@ -17,7 +17,7 @@ infer_device, infer_is_classifier, ) -from nvforest._handle import Handle +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], - handle: Optional[Handle], + stream: Optional[StreamLike], 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[StreamLike] = 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 : 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. """ 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[StreamLike] = 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 : 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. """ 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[StreamLike] = 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 : 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. """ 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..2d67e87 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 # @@ -18,8 +18,7 @@ from cuda.bindings import runtime from nvforest._base import ForestInferenceClassifier, ForestInferenceRegressor -from nvforest._handle import Handle -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, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -200,7 +199,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 +242,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 +259,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -272,7 +271,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 +283,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -296,7 +295,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 +396,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -409,7 +408,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 +420,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -433,7 +432,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 +525,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -539,7 +538,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 +550,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -563,7 +562,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 +664,7 @@ def __init__( self, *, treelite_model: treelite.Model, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -678,7 +677,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 +689,7 @@ def _create_with_layout( cls, *, treelite_model_bytes: bytes, - handle: Optional[Handle], + stream: Optional[StreamLike], layout: str, default_chunk_size: Optional[int], align_bytes: Optional[int], @@ -702,7 +701,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/_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/cuda_stream.pxd b/python/nvforest/nvforest/detail/cuda_stream.pxd index 690c0c4..0fea6ca 100644 --- a/python/nvforest/nvforest/detail/cuda_stream.pxd +++ b/python/nvforest/nvforest/detail/cuda_stream.pxd @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # +from cuda.bindings.cyruntime cimport cudaStream_t + + cdef extern from "nvforest/cuda_stream.hpp" namespace "nvforest" nogil: - cdef cppclass cuda_stream: - pass + 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..61aa29a 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 -from nvforest._handle import Handle -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 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,16 @@ 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() - ) + # 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, 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 cdef optional[bool] use_double_precision_c cdef bool use_double_precision_bool @@ -130,7 +132,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 +242,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 +275,7 @@ class ForestInferenceImpl: treelite_model: treelite.Model, device: str, device_id: int, - handle: Optional[Handle] = None, + stream: Optional[StreamLike] = None, layout: str = "depth_first", default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, @@ -281,7 +283,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, StreamLike): + raise TypeError("stream must be a stream-like object 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 +326,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+ From be09e5d8d0fdc054e8696738a64daf4aab8ff212 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 00:47:03 -0700 Subject: [PATCH 02/13] Undo unrelated change --- cpp/include/nvforest/forest_model.hpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index b504d6d..be305f4 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -19,7 +19,6 @@ #include #include -#include #include namespace nvforest { @@ -33,7 +32,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_{std::move(forest)} + : decision_forest_{forest} { } From 1f4a93d25df9cffa174d2ef702414db9fa69f7f9 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 09:32:05 -0700 Subject: [PATCH 03/13] Add missing cuda-core dep --- dependencies.yaml | 1 + python/nvforest/pyproject.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/dependencies.yaml b/dependencies.yaml index d85a743..13dd614 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -129,6 +129,7 @@ files: table: project includes: - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libnvforest - depends_on_pylibraft diff --git a/python/nvforest/pyproject.toml b/python/nvforest/pyproject.toml index 437fb9d..6a6bc43 100644 --- a/python/nvforest/pyproject.toml +++ b/python/nvforest/pyproject.toml @@ -28,6 +28,7 @@ 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", From fe009c8ab30a77927face7f251c0dde74d96477f Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 17:06:57 -0700 Subject: [PATCH 04/13] Make stream the first parameter --- cpp/include/nvforest/forest_model.hpp | 24 +++++++++---------- cpp/tests/treelite_importer.cpp | 8 +++---- .../nvforest/detail/forest_inference.pyx | 6 ++--- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index b504d6d..dea577e 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -116,12 +116,12 @@ struct forest_model { /** * Perform inference on given input * - * @param[out] output The buffer where model output should be stored. - * This must be of size at least ROWS x num_outputs(). - * @param[in] input The buffer containing input data. * @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. + * @param[out] output The buffer where model output should be stored. + * This must be of size at least ROWS x num_outputs(). + * @param[in] input The buffer containing input data. * @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. @@ -137,9 +137,9 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(buffer& output, + void predict(cuda_stream stream, + buffer& output, buffer const& input, - cuda_stream stream = cuda_stream{}, infer_kind predict_type = infer_kind::default_kind, std::optional specified_chunk_size = std::nullopt) { @@ -183,6 +183,10 @@ struct forest_model { /** * Perform inference on given input * + * @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[out] output Pointer to the memory location where output should end * up * @param[in] input Pointer to the input data @@ -190,10 +194,6 @@ 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. @@ -209,12 +209,12 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(io_t* output, + void predict(cuda_stream stream, + 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) { @@ -230,7 +230,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(out_buffer, in_buffer, stream, predict_type, specified_chunk_size); + predict(stream, out_buffer, in_buffer, predict_type, specified_chunk_size); } private: diff --git a/cpp/tests/treelite_importer.cpp b/cpp/tests/treelite_importer.cpp index 1b7f089..33403d0 100644 --- a/cpp/tests/treelite_importer.cpp +++ b/cpp/tests/treelite_importer.cpp @@ -385,12 +385,12 @@ TEST(TreeliteImporter, DegenerateTree) auto X = std::vector{0.0}; auto preds = std::vector(1, 0.0); auto expected_preds = std::vector{1.0}; - nvforest_model.predict(preds.data(), + nvforest_model.predict(nvforest::cuda_stream{}, + 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); @@ -405,12 +405,12 @@ TEST(TreeliteImporter, DegenerateTreeWithVectorLeaf) 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(preds.data(), + nvforest_model.predict(nvforest::cuda_stream{}, + 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); diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 61aa29a..203f2ee 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -31,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]( + nvforest_stream_t, io_t*, io_t*, size_t, nvforest_device_t, nvforest_device_t, - nvforest_stream_t, infer_kind, optional[uint32_t] ) except + @@ -242,23 +242,23 @@ cdef class ForestInference_impl(): if model_dtype == np.float32: self.model.predict[float]( + self.stream_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.stream_handle, out_ptr, in_ptr, n_rows, out_dev, in_dev, - self.stream_handle, infer_type_enum, chunk_specification ) From 9f8bf427a8c97f83ee5adfcbd1df0b230ead9021 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 17:23:15 -0700 Subject: [PATCH 05/13] Use cuda::stream_ref instead of cudaStream_t --- cpp/include/nvforest/buffer.hpp | 12 ++++++++---- cpp/include/nvforest/detail/owning_buffer/gpu.hpp | 5 +++-- cpp/src/detail/device_buffer.cu | 6 +++--- 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/cpp/include/nvforest/buffer.hpp b/cpp/include/nvforest/buffer.hpp index 267cad8..c2558fa 100644 --- a/cpp/include/nvforest/buffer.hpp +++ b/cpp/include/nvforest/buffer.hpp @@ -13,6 +13,8 @@ #include #include +#include + #include #include @@ -57,7 +59,8 @@ struct buffer { switch (mem_type) { case device_type::cpu: result = detail::owning_buffer{size}; break; case device_type::gpu: - result = detail::owning_buffer{std::get<1>(device_), size, stream}; + result = detail::owning_buffer{ + std::get<1>(device_), size, cuda::stream_ref{stream}}; break; } return result; @@ -138,8 +141,8 @@ struct buffer { result_data = buf.get(); result = std::move(buf); } else if (mem_type == device_type::gpu) { - auto buf = - detail::owning_buffer(std::get<1>(device_), other.size(), stream); + auto buf = detail::owning_buffer( + std::get<1>(device_), other.size(), cuda::stream_ref{stream}); result_data = buf.get(); result = std::move(buf); } @@ -212,7 +215,8 @@ struct buffer { result_data = buf.get(); result = std::move(buf); } else if (mem_type == device_type::gpu) { - auto buf = detail::owning_buffer{device, other.size(), stream}; + auto buf = detail::owning_buffer{ + device, other.size(), cuda::stream_ref{stream}}; result_data = buf.get(); result = std::move(buf); } diff --git a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp index 8cca3f4..e685578 100644 --- a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp +++ b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -21,7 +22,7 @@ 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); + cuda::stream_ref 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(); @@ -38,7 +39,7 @@ struct owning_buffer { owning_buffer() = default; owning_buffer(device_id device_id, std::size_t size, - cudaStream_t stream) noexcept(false) + cuda::stream_ref stream) noexcept(false) : data_{device_id, size * sizeof(value_type), stream} { } diff --git a/cpp/src/detail/device_buffer.cu b/cpp/src/detail/device_buffer.cu index a9fdf21..9271918 100644 --- a/cpp/src/detail/device_buffer.cu +++ b/cpp/src/detail/device_buffer.cu @@ -16,11 +16,11 @@ 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) + owning_device_buffer_type_erased_impl(int device_id, std::size_t size, cuda::stream_ref 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); + return cuda::make_buffer(stream, mr, size, cuda::no_init); }()} { } @@ -33,7 +33,7 @@ struct owning_device_buffer_type_erased_impl { 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) + device_id device_id, std::size_t size, cuda::stream_ref stream) { auto device_context = device_setter{device_id}; impl_ = std::make_unique(device_id.value(), size, stream); From d94a9c9b4af06adeeabd493f1a027537d589512f Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 17:26:55 -0700 Subject: [PATCH 06/13] Validate stream protocol version --- python/nvforest/nvforest/detail/forest_inference.pyx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 203f2ee..5180c0b 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -85,6 +85,8 @@ cdef class ForestInference_impl(): if not isinstance(stream, StreamLike): raise TypeError("stream must be a stream-like object or None") stream_tuple: tuple[int, int] = stream.__cuda_stream__() + if len(stream_tuple) != 2 or stream_tuple[0] != 0: + raise TypeError("stream must use the version 0 stream protocol") stream_ptr = stream_tuple[1] self.stream = stream self.stream_handle = stream_ptr From d6ea038631fc4361ee53f2d7626840e8faa60aef Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 17:55:01 -0700 Subject: [PATCH 07/13] Address coderabbit comments --- cpp/src/detail/device_buffer.cu | 2 +- .../nvforest/detail/forest_inference.pyx | 26 ++++++++++++------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/cpp/src/detail/device_buffer.cu b/cpp/src/detail/device_buffer.cu index 9271918..87fb5dd 100644 --- a/cpp/src/detail/device_buffer.cu +++ b/cpp/src/detail/device_buffer.cu @@ -45,6 +45,6 @@ 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(); } +std::byte* owning_device_buffer_type_erased::get() { return impl_ ? impl_->get() : nullptr; } } // namespace nvforest::detail diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 5180c0b..8d55a7a 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -7,7 +7,7 @@ from typing import Optional, Union import numpy as np import treelite -from cuda.core import Device +from cuda.core import Device, Stream from nvforest._typing import DataType, StreamLike from nvforest.detail.treelite import safe_treelite_call @@ -80,14 +80,17 @@ cdef class ForestInference_impl(): ): # Assumption: The caller needs to pass in correct (device, device_id) pair # This function will not contain any logic for auto-detecting device. + # Assumption: The caller should convert a user-provided stream-like object + # into cuda.core.Stream before constructing ForestInference_impl() cdef uintptr_t stream_ptr = 0 if stream is not None: - if not isinstance(stream, StreamLike): - raise TypeError("stream must be a stream-like object or None") - stream_tuple: tuple[int, int] = stream.__cuda_stream__() - if len(stream_tuple) != 2 or stream_tuple[0] != 0: - raise TypeError("stream must use the version 0 stream protocol") - stream_ptr = stream_tuple[1] + # Use assertion here, since the failure here indicates a bug, not + # a user error. + assert isinstance(stream, Stream), ( + "stream must be converted to cuda.core.Stream before " + "building ForestInference_impl()" + ) + stream_ptr = int(stream.handle) self.stream = stream self.stream_handle = stream_ptr @@ -287,14 +290,19 @@ class ForestInferenceImpl: # This function will not contain any logic for auto-detecting device. 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: + if device == "gpu": previous_device = Device() try: cuda_device = Device(device_id) cuda_device.set_current() - stream = cuda_device.create_stream() + if stream is None: + stream = cuda_device.create_stream() + else: + stream = cuda_device.create_stream(stream) finally: previous_device.set_current() + else: + assert device == "cpu" if device == "gpu" and stream.device.device_id != device_id: raise ValueError( f"stream is associated with device {stream.device.device_id}, " From d9e15b2e1b91bc9b889317a1133e3447dcd2c8ed Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 18:26:26 -0700 Subject: [PATCH 08/13] Ensure that stream is on the correct device --- cpp/include/nvforest/detail/infer/gpu.cuh | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/cpp/include/nvforest/detail/infer/gpu.cuh b/cpp/include/nvforest/detail/infer/gpu.cuh index 744f48c..22c1dea 100644 --- a/cpp/include/nvforest/detail/infer/gpu.cuh +++ b/cpp/include/nvforest/detail/infer/gpu.cuh @@ -21,6 +21,8 @@ #include #include +#include + #include #include #include @@ -41,6 +43,20 @@ inline auto compute_output_size(index_type row_output_size, return result; } +// If a non-default stream is provided, it must reside on the correct device. +inline void validate_stream(device_id device, cuda_stream stream) +{ + cuda::stream_ref legacy_default_stream{cudaStreamLegacy}; + cuda::stream_ref per_thread_default_stream{cudaStreamPerThread}; + cuda::stream_ref stream_wrapped{stream}; + if (stream_wrapped != legacy_default_stream && stream_wrapped != per_thread_default_stream && + stream_wrapped.device().get() != device.value()) { + throw std::runtime_error{std::string("Stream on the wrong device. ") + + "Expected: " + std::to_string(device.value()) + + ", Actual: " + std::to_string(stream_wrapped.device().get())}; + } +} + /* A wrapper around the underlying inference kernels to support dispatching to * the right kernel * @@ -99,6 +115,8 @@ std::enable_if_t infer( { using output_t = typename forest_t::template raw_output_type; + validate_stream(device, stream); + auto sm_count = get_sm_count(device); auto const max_shared_mem_per_block = get_max_shared_mem_per_block(device); auto const max_shared_mem_per_sm = get_max_shared_mem_per_sm(device); From 72b8d885f0d7bc27b40ce3c9f49e0370bb3bbcff Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 18:42:02 -0700 Subject: [PATCH 09/13] Use cuda::stream_ref in validate_stream() --- cpp/include/nvforest/detail/infer/gpu.cuh | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/cpp/include/nvforest/detail/infer/gpu.cuh b/cpp/include/nvforest/detail/infer/gpu.cuh index 22c1dea..31cc8ea 100644 --- a/cpp/include/nvforest/detail/infer/gpu.cuh +++ b/cpp/include/nvforest/detail/infer/gpu.cuh @@ -44,16 +44,15 @@ inline auto compute_output_size(index_type row_output_size, } // If a non-default stream is provided, it must reside on the correct device. -inline void validate_stream(device_id device, cuda_stream stream) +inline void validate_stream(device_id device, cuda::stream_ref stream) { cuda::stream_ref legacy_default_stream{cudaStreamLegacy}; cuda::stream_ref per_thread_default_stream{cudaStreamPerThread}; - cuda::stream_ref stream_wrapped{stream}; - if (stream_wrapped != legacy_default_stream && stream_wrapped != per_thread_default_stream && - stream_wrapped.device().get() != device.value()) { + if (stream != legacy_default_stream && stream != per_thread_default_stream && + stream.device().get() != device.value()) { throw std::runtime_error{std::string("Stream on the wrong device. ") + "Expected: " + std::to_string(device.value()) + - ", Actual: " + std::to_string(stream_wrapped.device().get())}; + ", Actual: " + std::to_string(stream.device().get())}; } } @@ -115,7 +114,7 @@ std::enable_if_t infer( { using output_t = typename forest_t::template raw_output_type; - validate_stream(device, stream); + validate_stream(device, cuda::stream_ref{stream}); auto sm_count = get_sm_count(device); auto const max_shared_mem_per_block = get_max_shared_mem_per_block(device); From 37d1da8773cc08c68bf659af32095ab3619c7697 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 19:00:24 -0700 Subject: [PATCH 10/13] Add test coverage for streams on wrong device --- cpp/tests/CMakeLists.txt | 1 + cpp/tests/invalid_stream.cu | 82 ++++++++++++++++++++++++++ python/nvforest/tests/test_nvforest.py | 34 +++++++++++ 3 files changed, 117 insertions(+) create mode 100644 cpp/tests/invalid_stream.cu diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index d4aeed8..7046f9b 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -81,6 +81,7 @@ endfunction() ConfigureTest(NAME HOST_BUFFER_TEST buffer/buffer.cpp) if(NVFOREST_ENABLE_GPU) ConfigureTest(NAME DEVICE_BUFFER_TEST buffer/buffer.cu) + ConfigureTest(NAME INVALID_STREAM_TEST invalid_stream.cu) endif() ConfigureTest(NAME FOREST_TRAVERSAL_TEST forest/traversal_forest.cpp) ConfigureTest(NAME TREELITE_TRAVERSAL_TEST forest/treelite_traversal.cpp) diff --git a/cpp/tests/invalid_stream.cu b/cpp/tests/invalid_stream.cu new file mode 100644 index 0000000..103b2b4 --- /dev/null +++ b/cpp/tests/invalid_stream.cu @@ -0,0 +1,82 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +#include +#include + +namespace nvforest { + +TEST(InvalidStream, wrong_device) +{ + auto const num_devices = cuda::devices.size(); + if (num_devices < 2) { GTEST_SKIP() << "Test requires at least 2 GPU devices"; } + + // Make a stump tree + auto model_builder = + treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32, + treelite::TypeInfo::kFloat32, + treelite::model_builder::Metadata{ + 1, + treelite::TaskType::kRegressor, + false, + 1, + {1}, + {1, 1}, + }, + treelite::model_builder::TreeAnnotation{1, {0}, {0}}, + treelite::model_builder::PostProcessorFunc{"identity"}, + std::vector{0.0}); + model_builder->StartTree(); + model_builder->StartNode(0); + model_builder->NumericalTest(0, 0.0, true, treelite::Operator::kLT, 1, 2); + model_builder->EndNode(); + model_builder->StartNode(1); + model_builder->LeafScalar(-1.0); + model_builder->EndNode(); + model_builder->StartNode(2); + model_builder->LeafScalar(1.0); + model_builder->EndNode(); + model_builder->EndTree(); + + cuda_stream stream; + { + auto device_context = detail::device_setter{detail::device_id{1}}; + detail::cuda_check(cudaStreamCreate(&stream)); + } + + // Loading a model on device 0 and inferencing on device 1 is a ner + auto tl_model = model_builder->CommitModel(); + auto nvforest_model = import_from_treelite_model( + *tl_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu, 0, stream); + auto input = thrust::device_vector{1.0f}; + auto output = thrust::device_vector(1); + EXPECT_THAT( + [&]() { + nvforest_model.predict(stream, + thrust::raw_pointer_cast(output.data()), + thrust::raw_pointer_cast(input.data()), + 1, + device_type::gpu, + device_type::gpu, + infer_kind::default_kind); + }, + testing::ThrowsMessage( + testing::HasSubstr("Stream on the wrong device. Expected: 0, Actual: 1"))); +} + +} // namespace nvforest diff --git a/python/nvforest/tests/test_nvforest.py b/python/nvforest/tests/test_nvforest.py index a7282ae..1d5bb4e 100644 --- a/python/nvforest/tests/test_nvforest.py +++ b/python/nvforest/tests/test_nvforest.py @@ -927,3 +927,37 @@ def test_incorrect_data_shape(input_size, predict_func): with pytest.raises(ValueError, match=f"Expected {n_features} features"): X_test = np.zeros((1, input_size)) _ = predict_func(fm, X_test) + + +@pytest.mark.skipif( + cp.cuda.runtime.getDeviceCount() < 2, reason="Requires at least 2 GPUs" +) +def test_incompatible_stream(tmp_path): + random_state = np.random.RandomState(43210) + + X, y = _simulate_data( + 1000, + 2, + random_state=random_state, + classification=False, + bias=10.0, + ) + + model_path = tmp_path / "xgb_class.ubj" + _ = _build_and_save_xgboost( + model_path, + X, + y, + num_rounds=10, + classification=False, + ) + with cp.cuda.Device(1): + stream = cp.cuda.Stream(non_blocking=True) + assert stream.device_id == 1 + with pytest.raises( + ValueError, + match=r".*stream is associated with device 1, but device_id is 0.*", + ): + _ = nvforest.load_model( + model_path, device="gpu", device_id=0, stream=stream + ) From 4276fb3021856afa65968b7f7b5e7c0848cb3ad9 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 19:35:00 -0700 Subject: [PATCH 11/13] Fix CPU build --- cpp/include/nvforest/buffer.hpp | 12 ++++-------- cpp/include/nvforest/detail/owning_buffer/gpu.hpp | 5 +++-- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/cpp/include/nvforest/buffer.hpp b/cpp/include/nvforest/buffer.hpp index c2558fa..267cad8 100644 --- a/cpp/include/nvforest/buffer.hpp +++ b/cpp/include/nvforest/buffer.hpp @@ -13,8 +13,6 @@ #include #include -#include - #include #include @@ -59,8 +57,7 @@ struct buffer { switch (mem_type) { case device_type::cpu: result = detail::owning_buffer{size}; break; case device_type::gpu: - result = detail::owning_buffer{ - std::get<1>(device_), size, cuda::stream_ref{stream}}; + result = detail::owning_buffer{std::get<1>(device_), size, stream}; break; } return result; @@ -141,8 +138,8 @@ struct buffer { result_data = buf.get(); result = std::move(buf); } else if (mem_type == device_type::gpu) { - auto buf = detail::owning_buffer( - std::get<1>(device_), other.size(), cuda::stream_ref{stream}); + auto buf = + detail::owning_buffer(std::get<1>(device_), other.size(), stream); result_data = buf.get(); result = std::move(buf); } @@ -215,8 +212,7 @@ struct buffer { result_data = buf.get(); result = std::move(buf); } else if (mem_type == device_type::gpu) { - auto buf = detail::owning_buffer{ - device, other.size(), cuda::stream_ref{stream}}; + auto buf = detail::owning_buffer{device, other.size(), stream}; result_data = buf.get(); result = std::move(buf); } diff --git a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp index e685578..2e217d6 100644 --- a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp +++ b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include #include #include #include @@ -39,8 +40,8 @@ struct owning_buffer { owning_buffer() = default; owning_buffer(device_id device_id, std::size_t size, - cuda::stream_ref stream) noexcept(false) - : data_{device_id, size * sizeof(value_type), stream} + cuda_stream stream) noexcept(false) + : data_{device_id, size * sizeof(value_type), cuda::stream_ref{stream}} { } From a483b52605496975a924c596c552e18f945129cc Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 13 Aug 2026 19:35:07 -0700 Subject: [PATCH 12/13] Fix typo --- cpp/tests/invalid_stream.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/invalid_stream.cu b/cpp/tests/invalid_stream.cu index 103b2b4..0d9b028 100644 --- a/cpp/tests/invalid_stream.cu +++ b/cpp/tests/invalid_stream.cu @@ -59,7 +59,7 @@ TEST(InvalidStream, wrong_device) detail::cuda_check(cudaStreamCreate(&stream)); } - // Loading a model on device 0 and inferencing on device 1 is a ner + // Loading a model on device 0 and inferencing on device 1 is an error. auto tl_model = model_builder->CommitModel(); auto nvforest_model = import_from_treelite_model( *tl_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu, 0, stream); From fc24a343c969021e455e13e15b5ba13b24d38c6e Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 20 Aug 2026 22:55:38 -0700 Subject: [PATCH 13/13] Add back handle arg + mark it deprecated --- cpp/CMakeLists.txt | 5 +- cpp/cmake/thirdparty/get_raft.cmake | 78 +++++++++++++++++++ python/nvforest/CMakeLists.txt | 6 ++ python/nvforest/nvforest/_factory.py | 23 ++++++ python/nvforest/nvforest/_forest_inference.py | 11 +++ python/nvforest/nvforest/_handle.py | 9 +++ .../nvforest/nvforest/detail/raft_stream.pyx | 35 +++++++++ 7 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 cpp/cmake/thirdparty/get_raft.cmake create mode 100644 python/nvforest/nvforest/_handle.py create mode 100644 python/nvforest/nvforest/detail/raft_stream.pyx diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index c8d4113..d0f8d2b 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -198,6 +198,9 @@ endif() # CCCL before RMM, and RMM before RAFT include(cmake/thirdparty/get_cccl.cmake) +if(NVFOREST_ENABLE_GPU) + include(cmake/thirdparty/get_raft.cmake) +endif() include(cmake/thirdparty/get_treelite.cmake) if(BUILD_NVFOREST_TESTS) @@ -309,7 +312,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE) endif() if(NVFOREST_ENABLE_GPU) - list(APPEND _nvforest_cpp_public_libs CUDA::cudart_static CCCL::CCCL) + list(APPEND _nvforest_cpp_public_libs raft::raft CUDA::cudart_static CCCL::CCCL) endif() # These are always private: 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/python/nvforest/CMakeLists.txt b/python/nvforest/CMakeLists.txt index 5ebd872..ce399e6 100644 --- a/python/nvforest/CMakeLists.txt +++ b/python/nvforest/CMakeLists.txt @@ -77,3 +77,9 @@ rapids_cython_create_modules( SOURCE_FILES "nvforest/detail/treelite.pyx" LINKED_LIBRARIES "${linked_libraries}" INSTALL_DIR nvforest/detail MODULE_PREFIX nvforest_) + +rapids_cython_create_modules( + CXX + SOURCE_FILES "nvforest/detail/raft_stream.pyx" + LINKED_LIBRARIES "${linked_libraries}" + INSTALL_DIR nvforest/detail MODULE_PREFIX nvforest_) diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index b999917..f541822 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -17,7 +17,9 @@ infer_device, infer_is_classifier, ) +from nvforest._handle import Handle from nvforest._typing import StreamLike +from nvforest.detail.raft_stream import _handle_deprecated_handle_arg def get_forest_inference_class(device, is_classifier) -> type: @@ -73,6 +75,7 @@ def load_model( precision: Optional[str] = None, device_id: Optional[int] = None, stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, ) -> ForestInference: """Load a model into nvForest from a serialized model file. @@ -117,6 +120,11 @@ def load_model( 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. + handle : nvforest.Handle or None + **Deprecated**. This argument will be removed in 26.12. Please use + the ``stream`` argument instead. + For GPU execution, the nvForest handle containing the stream or stream + pool to use during loading and inference. """ model_path = pathlib.Path(model_file) if not model_path.exists(): @@ -153,6 +161,7 @@ def load_model( case _: raise ValueError(f"Unknown model type: {model_type}") + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, @@ -175,6 +184,7 @@ def load_from_sklearn( precision: Optional[str] = None, device_id: Optional[int] = None, stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, ) -> ForestInference: """Load a Scikit-Learn forest model to nvForest @@ -212,9 +222,15 @@ def load_from_sklearn( 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. + handle : nvforest.Handle or None + **Deprecated**. This argument will be removed in 26.12. Please use + the ``stream`` argument instead. + For GPU execution, the nvForest handle containing the stream or stream + pool to use during loading and inference. """ tl_model = treelite.sklearn.import_model(skl_model) + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, @@ -237,6 +253,7 @@ def load_from_treelite_model( precision: Optional[str] = None, device_id: Optional[int] = None, stream: Optional[StreamLike] = None, + handle: Optional[Handle] = None, ) -> ForestInference: """Load a Treelite forest model to nvForest @@ -274,7 +291,13 @@ def load_from_treelite_model( 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. + handle : nvforest.Handle or None + **Deprecated**. This argument will be removed in 26.12. Please use + the ``stream`` argument instead. + For GPU execution, the nvForest handle containing the stream or stream + pool to use during loading and inference. """ + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) return make_forest_inference_object( treelite_model=tl_model, device=device, diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index 2d67e87..8885275 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -8,6 +8,9 @@ from time import perf_counter from typing import Optional +from nvforest._handle import Handle +from nvforest.detail.raft_stream import _handle_deprecated_handle_arg + try: from typing import Self except ImportError: @@ -264,9 +267,11 @@ def __init__( default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, + handle: Optional[Handle] = None, ): if not infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a classifier.") + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="cpu", @@ -401,9 +406,11 @@ def __init__( default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, + handle: Optional[Handle] = None, ): if infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a regressor.") + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="cpu", @@ -530,10 +537,12 @@ def __init__( default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, + handle: Optional[Handle] = None, device_id: int, ): if not infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a classifier.") + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="gpu", @@ -669,10 +678,12 @@ def __init__( default_chunk_size: Optional[int] = None, align_bytes: Optional[int] = None, precision: Optional[str] = None, + handle: Optional[Handle] = None, device_id: int, ): if infer_is_classifier(treelite_model): raise ValueError("treelite_model must be a regressor.") + stream = _handle_deprecated_handle_arg(handle=handle, stream=stream) self.forest = ForestInferenceImpl( treelite_model=treelite_model, device="gpu", 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 diff --git a/python/nvforest/nvforest/detail/raft_stream.pyx b/python/nvforest/nvforest/detail/raft_stream.pyx new file mode 100644 index 0000000..e47be55 --- /dev/null +++ b/python/nvforest/nvforest/detail/raft_stream.pyx @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import warnings +from typing import Optional + +from cuda.core import Stream +from pylibraft.common.handle import Handle as RaftHandle + +from nvforest._typing import StreamLike + + +def _get_stream_from_raft_handle(handle: RaftHandle) -> Stream: + cdef int stream_ptr = handle.c_obj.get_stream() + return Stream.from_handle(stream_ptr) + + +def _handle_deprecated_handle_arg( + *, + handle: Optional[RaftHandle] = None, + stream: Optional[StreamLike] = None, +) -> Optional[StreamLike]: + if handle is None: + return stream + + if stream is not None: + raise ValueError("Cannot set `stream` and `handle` at the same time.") + warnings.warn( + "`handle` parameter is deprecated and will be removed in 26.12. " + "Please use `stream` instead", + FutureWarning, + stacklevel=2, + ) + stream = _get_stream_from_raft_handle(handle) + return stream