From 1a988ddcf076a0b17f5c3a78b6fc1cb2313322e6 Mon Sep 17 00:00:00 2001 From: Philip Hyunsu Cho Date: Tue, 8 Sep 2026 06:27:51 -0700 Subject: [PATCH 01/20] Remove RMM dependency and prepare for RAFT removal, 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 | 6 +- cpp/cmake/thirdparty/get_rmm.cmake | 13 -- cpp/include/nvforest/detail/device_id/gpu.hpp | 11 +- .../nvforest/detail/device_setter/gpu.hpp | 18 ++- cpp/include/nvforest/detail/infer/gpu.cuh | 17 +++ .../nvforest/detail/owning_buffer/gpu.hpp | 38 +++-- cpp/include/nvforest/forest_model.hpp | 135 ++++-------------- cpp/include/nvforest/handle.hpp | 43 ------ cpp/include/nvforest/treelite_importer.hpp | 1 - cpp/src/detail/device_buffer.cu | 50 +++++++ cpp/tests/CMakeLists.txt | 5 +- cpp/tests/invalid_stream.cu | 82 +++++++++++ cpp/tests/treelite_importer.cpp | 18 +-- dependencies.yaml | 8 ++ python/nvforest/CMakeLists.txt | 6 + python/nvforest/nvforest/__init__.py | 4 +- python/nvforest/nvforest/_factory.py | 45 ++++-- python/nvforest/nvforest/_forest_inference.py | 52 ++++--- python/nvforest/nvforest/_typing.py | 9 +- .../nvforest/nvforest/detail/cuda_stream.pxd | 6 +- .../nvforest/detail/forest_inference.pyx | 69 ++++++--- python/nvforest/nvforest/detail/handle.pxd | 17 --- .../nvforest/nvforest/detail/raft_stream.pyx | 35 +++++ python/nvforest/pyproject.toml | 1 + python/nvforest/tests/test_nvforest.py | 34 +++++ 29 files changed, 447 insertions(+), 280 deletions(-) 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 create mode 100644 cpp/tests/invalid_stream.cu delete mode 100644 python/nvforest/nvforest/detail/handle.pxd create mode 100644 python/nvforest/nvforest/detail/raft_stream.pyx diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 26621dfb..734d3ff1 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 0a8e2725..ac607abc 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 48bc9bad..ac0b0b6b 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 4678e524..3b9c47cf 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 885fa544..d0f8d2b2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -199,7 +199,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) @@ -252,7 +251,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 +312,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 raft::raft CUDA::cudart_static CCCL::CCCL) endif() # These are always private: diff --git a/cpp/cmake/thirdparty/get_rmm.cmake b/cpp/cmake/thirdparty/get_rmm.cmake deleted file mode 100644 index 519b9905..00000000 --- a/cpp/cmake/thirdparty/get_rmm.cmake +++ /dev/null @@ -1,13 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# 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 49edfdd6..10d10119 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 d69a0ff5..07ba9a34 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/infer/gpu.cuh b/cpp/include/nvforest/detail/infer/gpu.cuh index 744f48cf..31cc8eaa 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,19 @@ 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_ref stream) +{ + cuda::stream_ref legacy_default_stream{cudaStreamLegacy}; + cuda::stream_ref per_thread_default_stream{cudaStreamPerThread}; + 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.device().get())}; + } +} + /* A wrapper around the underlying inference kernels to support dispatching to * the right kernel * @@ -99,6 +114,8 @@ std::enable_if_t infer( { using output_t = typename forest_t::template raw_output_type; + 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); auto const max_shared_mem_per_sm = get_max_shared_mem_per_sm(device); diff --git a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp index ba6030ff..2e217d64 100644 --- a/cpp/include/nvforest/detail/owning_buffer/gpu.hpp +++ b/cpp/include/nvforest/detail/owning_buffer/gpu.hpp @@ -3,37 +3,51 @@ * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include #include -#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, + 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(); + 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}}; - }()} + cuda_stream stream) noexcept(false) + : data_{device_id, size * sizeof(value_type), cuda::stream_ref{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 143881f4..dea577e5 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 @@ -118,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. @@ -139,113 +137,40 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(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) - { - 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, + void predict(cuda_stream stream, 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]( + [this, predict_type, &output, &input, &stream, &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 { @@ -258,8 +183,10 @@ 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[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 @@ -282,7 +209,7 @@ struct forest_model { * reasonable value. On CPU, this argument can generally just be omitted. */ template - void predict(handle_t const& handle, + void predict(cuda_stream stream, io_t* output, io_t* input, std::size_t num_rows, @@ -303,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(stream, out_buffer, in_buffer, 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 fd1f3bb7..00000000 --- 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 5402b880..b504403f 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 00000000..87fb5ddd --- /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, 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(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, cuda::stream_ref 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_ ? impl_->get() : nullptr; } + +} // namespace nvforest::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 96f5297e..7046f9ba 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}>") @@ -85,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 00000000..0d9b028b --- /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 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); + 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/cpp/tests/treelite_importer.cpp b/cpp/tests/treelite_importer.cpp index 37d3d9c8..33403d05 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,16 +382,10 @@ 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, + nvforest_model.predict(nvforest::cuda_stream{}, preds.data(), X.data(), 1, @@ -408,16 +402,10 @@ 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, + nvforest_model.predict(nvforest::cuda_stream{}, preds.data(), X.data(), 1, diff --git a/dependencies.yaml b/dependencies.yaml index b0e1d991..63478f73 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 @@ -127,6 +129,7 @@ files: table: project includes: - depends_on_cuda_python + - depends_on_cuda_core - depends_on_cupy - depends_on_libnvforest - depends_on_pylibraft @@ -455,6 +458,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/CMakeLists.txt b/python/nvforest/CMakeLists.txt index 5ebd872c..ce399e66 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/__init__.py b/python/nvforest/nvforest/__init__.py index f92070eb..e5a01302 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -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 96b44546..0b1091f1 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -18,6 +18,8 @@ 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: @@ -39,7 +41,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 +52,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,6 +74,7 @@ def load_model( align_bytes: Optional[int] = None, 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. @@ -113,10 +116,15 @@ 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 : 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. handle : nvforest.Handle or None + **Deprecated**. This argument will be removed in 27.02. Please use + the ``stream`` argument instead. 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. + pool to use during loading and inference. """ model_path = pathlib.Path(model_file) if not model_path.exists(): @@ -153,11 +161,12 @@ 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, device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -174,6 +183,7 @@ def load_from_sklearn( align_bytes: Optional[int] = None, 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 @@ -208,18 +218,24 @@ 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 : 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. handle : nvforest.Handle or None + **Deprecated**. This argument will be removed in 27.02. Please use + the ``stream`` argument instead. 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. + 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, device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -236,6 +252,7 @@ def load_from_treelite_model( align_bytes: Optional[int] = None, 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 @@ -270,16 +287,22 @@ 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 : 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. handle : nvforest.Handle or None + **Deprecated**. This argument will be removed in 27.02. Please use + the ``stream`` argument instead. 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. + 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, 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 c892d50b..88852756 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: @@ -18,8 +21,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 +125,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 +202,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 +245,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,19 +262,21 @@ 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, 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", device_id=-1, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -284,7 +288,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 +300,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,19 +401,21 @@ 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, 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", device_id=-1, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -421,7 +427,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 +439,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,20 +532,22 @@ 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, 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", device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -551,7 +559,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 +571,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,20 +673,22 @@ 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, 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", device_id=device_id, - handle=handle, + stream=stream, layout=layout, default_chunk_size=default_chunk_size, align_bytes=align_bytes, @@ -690,7 +700,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 +712,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/_typing.py b/python/nvforest/nvforest/_typing.py index 41c89160..5f897fbd 100644 --- a/python/nvforest/nvforest/_typing.py +++ b/python/nvforest/nvforest/_typing.py @@ -3,7 +3,7 @@ # 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 690c0c4c..0fea6cab 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 3d6eb3b4..8d55a7a7 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._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,7 +31,7 @@ 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&, + nvforest_stream_t, io_t*, io_t*, size_t, @@ -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,21 @@ 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. + # 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: + # 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 cdef optional[bool] use_double_precision_c cdef bool use_double_precision_bool @@ -130,7 +137,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,7 +247,7 @@ cdef class ForestInference_impl(): if model_dtype == np.float32: self.model.predict[float]( - self.nvforest_handle, + self.stream_handle, out_ptr, in_ptr, n_rows, @@ -251,7 +258,7 @@ cdef class ForestInference_impl(): ) else: self.model.predict[double]( - self.nvforest_handle, + self.stream_handle, out_ptr, in_ptr, n_rows, @@ -262,7 +269,7 @@ cdef class ForestInference_impl(): ) if self.device == "gpu": - self.nvforest_handle.synchronize() + self.stream.sync() return preds @@ -273,7 +280,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 +288,27 @@ 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": + previous_device = Device() + try: + cuda_device = Device(device_id) + cuda_device.set_current() + 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}, " + 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 +336,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 b80ad6b0..00000000 --- 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/nvforest/detail/raft_stream.pyx b/python/nvforest/nvforest/detail/raft_stream.pyx new file mode 100644 index 00000000..9c0376cb --- /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 27.02. " + "Please use `stream` instead", + FutureWarning, + stacklevel=2, + ) + stream = _get_stream_from_raft_handle(handle) + return stream diff --git a/python/nvforest/pyproject.toml b/python/nvforest/pyproject.toml index 6d8113f2..ed71be10 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.12.*,>=0.0.0a0", "numpy>=2.0,<3.0a0", diff --git a/python/nvforest/tests/test_nvforest.py b/python/nvforest/tests/test_nvforest.py index a7282ae4..1d5bb4e2 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 20264c8f6b6f37042bed718fd2fb231de161863a Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 10 Sep 2026 18:38:38 -0700 Subject: [PATCH 02/20] Add back Handle; fix _get_stream_from_raft_handle() --- python/nvforest/nvforest/__init__.py | 2 + .../nvforest/nvforest/detail/raft_stream.pyx | 11 ++- python/nvforest/tests/test_compat.py | 68 +++++++++++++++++++ 3 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 python/nvforest/tests/test_compat.py diff --git a/python/nvforest/nvforest/__init__.py b/python/nvforest/nvforest/__init__.py index e5a01302..1bd8d2f9 100644 --- a/python/nvforest/nvforest/__init__.py +++ b/python/nvforest/nvforest/__init__.py @@ -24,6 +24,7 @@ GPUForestInferenceClassifier, GPUForestInferenceRegressor, ) +from nvforest._handle import Handle from nvforest._typing import StreamLike from nvforest._version import __git_commit__, __version__ @@ -32,6 +33,7 @@ "CPUForestInferenceRegressor", "GPUForestInferenceClassifier", "GPUForestInferenceRegressor", + "Handle", "StreamLike", "load_model", "load_from_sklearn", diff --git a/python/nvforest/nvforest/detail/raft_stream.pyx b/python/nvforest/nvforest/detail/raft_stream.pyx index 9c0376cb..a90de6c2 100644 --- a/python/nvforest/nvforest/detail/raft_stream.pyx +++ b/python/nvforest/nvforest/detail/raft_stream.pyx @@ -9,10 +9,15 @@ from pylibraft.common.handle import Handle as RaftHandle from nvforest._typing import StreamLike +from cuda.bindings.cyruntime cimport cudaStream_t +from libc.stdint cimport uintptr_t +from pylibraft.common.handle cimport handle_t + def _get_stream_from_raft_handle(handle: RaftHandle) -> Stream: - cdef int stream_ptr = handle.c_obj.get_stream() - return Stream.from_handle(stream_ptr) + cdef handle_t* handle_ = handle.getHandle() + cdef cudaStream_t stream = handle_.get_stream() + return Stream.from_handle(stream) def _handle_deprecated_handle_arg( @@ -27,7 +32,7 @@ def _handle_deprecated_handle_arg( raise ValueError("Cannot set `stream` and `handle` at the same time.") warnings.warn( "`handle` parameter is deprecated and will be removed in 27.02. " - "Please use `stream` instead", + "Please use `stream` instead.", FutureWarning, stacklevel=2, ) diff --git a/python/nvforest/tests/test_compat.py b/python/nvforest/tests/test_compat.py new file mode 100644 index 00000000..4d068a33 --- /dev/null +++ b/python/nvforest/tests/test_compat.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Compatibility test to ensure that legacy APIs are working +""" + +import pytest +import treelite +from cuda.core import Device + +import nvforest + + +def test_deprecated_handle(): + """Test use of RAFT handle with non-default stream""" + builder = treelite.model_builder.ModelBuilder( + threshold_type="float64", + leaf_output_type="float64", + metadata=treelite.model_builder.Metadata( + num_feature=1, + task_type="kRegressor", + average_tree_output=False, + num_target=1, + num_class=[1], + leaf_vector_shape=(1, 1), + ), + tree_annotation=treelite.model_builder.TreeAnnotation( + num_tree=1, target_id=[0], class_id=[0] + ), + postprocessor=treelite.model_builder.PostProcessorFunc( + name="identity" + ), + base_scores=[0.0], + ) + builder.start_tree() + builder.start_node(0) + builder.numerical_test( + feature_id=0, + threshold=0.0, + default_left=False, + opname="<", + left_child_key=1, + right_child_key=2, + ) + builder.end_node() + builder.start_node(1) + builder.leaf(-1.0) + builder.end_node() + builder.start_node(2) + builder.leaf(1.0) + builder.end_node() + builder.end_tree() + + treelite_model = builder.commit() + + device = Device() + device.set_current() + stream = device.create_stream() + handle = nvforest.Handle(stream=stream) + + with pytest.warns( + FutureWarning, + match=".*`handle` parameter is deprecated and will be removed in 27.02.*", + ): + _ = nvforest.load_from_treelite_model( + treelite_model, device="gpu", handle=handle + ) From 626a2f31518673d031a77a9cc945e80175e514fe Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 10 Sep 2026 18:40:33 -0700 Subject: [PATCH 03/20] Set stream=None for CPU inference path --- python/nvforest/nvforest/detail/forest_inference.pyx | 1 + 1 file changed, 1 insertion(+) diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 8d55a7a7..295722eb 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -303,6 +303,7 @@ class ForestInferenceImpl: previous_device.set_current() else: assert device == "cpu" + stream = None if device == "gpu" and stream.device.device_id != device_id: raise ValueError( f"stream is associated with device {stream.device.device_id}, " From 0ea6b71db117ac59437368dccf57eb1f0f5e5359 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 10 Sep 2026 18:56:32 -0700 Subject: [PATCH 04/20] Update the examples --- README.md | 24 ++++++++++++++ cpp/include/nvforest/README.md | 20 ++++++------ docs/source/getting_started.rst | 58 ++++++++++++++++++++++++++++----- 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 11830123..aa102907 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,30 @@ fm = nvforest.load_model("/path/to/xgboost_model.ubj", device="gpu", pred = fm.predict(X) ``` +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 + +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) # Note: automatically syncs the stream +``` + Load a scikit-learn random forest model and get class probabilities: ```python diff --git a/cpp/include/nvforest/README.md b/cpp/include/nvforest/README.md index 711ab330..d6f028ef 100644 --- a/cpp/include/nvforest/README.md +++ b/cpp/include/nvforest/README.md @@ -98,25 +98,25 @@ auto output = static_cast(nullptr); // Loaded as single 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{}; - +auto stream = cudaStream_t{}; nvforest_model.predict( - handle, + stream, output, input, num_rows, nvforest::device_type::gpu, // out_mem_type nvforest::device_type::gpu, // in_mem_type + 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. +**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. **output**: Pointer to pre-allocated buffer where results should be written. If the model has been loaded at single precision, this should be a @@ -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/getting_started.rst b/docs/source/getting_started.rst index 8d48574a..12b2f465 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -128,6 +128,33 @@ 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 +--------------------- + +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 + + 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. + +.. [#] A stream-like object is an object that exposes a method named `__cuda_stream__`. + See `Stream Protocol `_ for more details. + Running inference ----------------- @@ -222,29 +249,44 @@ 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(stream, output, input, num_rows, nvforest::device_type::gpu, nvforest::device_type::gpu, nvforest::infer_kind::default_kind); + + // Wait before reading or freeing output. + cudaStreamSynchronize(stream); + + // Use output + // ... + + // Clean up + cudaStreamDestroy(stream); + // ... also free input, output buffers + +Use a stream associated with the same device as the model. Re-use the stream for import and +inference to preserve the ordering between model initialization and prediction without a global +synchronization. Keep the stream alive until the model is destroyed and all queued +work has completed. From 8c23909d935486d5a3a21cd49ef0dc84436238c8 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Thu, 10 Sep 2026 19:28:50 -0700 Subject: [PATCH 05/20] Update docs --- cpp/include/nvforest/README.md | 1 - docs/source/getting_started.rst | 9 +++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/include/nvforest/README.md b/cpp/include/nvforest/README.md index d6f028ef..bb3af152 100644 --- a/cpp/include/nvforest/README.md +++ b/cpp/include/nvforest/README.md @@ -98,7 +98,6 @@ auto output = static_cast(nullptr); // Loaded as single cudaMalloc((void**)&output, num_rows * num_outputs * sizeof(float)); // Assuming that input is a float* pointing to data already located on-device -auto stream = cudaStream_t{}; nvforest_model.predict( stream, output, diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 12b2f465..ff926557 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -252,6 +252,8 @@ Once the tree model is available as a Treelite object, pass it to the #include #include + constexpr int device_id = 0; + cudaSetDevice(device_id); cudaStream_t stream{}; cudaStreamCreate(&stream); @@ -261,7 +263,7 @@ Once the tree model is available as a Treelite object, pass it to the nvforest::index_type{}, std::nullopt, nvforest::device_type::gpu, - 0, + device_id, stream); Now that the tree model is fully imported into nvForest, let's run inference: @@ -282,9 +284,8 @@ Now that the tree model is fully imported into nvForest, let's run inference: // Use output // ... - // Clean up - cudaStreamDestroy(stream); - // ... also free input, output buffers + // Clean up by destroying stream + // Also free input, output buffers Use a stream associated with the same device as the model. Re-use the stream for import and inference to preserve the ordering between model initialization and prediction without a global From 04284878d7f50a69acd10cd3e3c967aa1d989882 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 14 Sep 2026 20:38:22 -0700 Subject: [PATCH 06/20] Fix hash for actions/checkout --- .github/workflows/link-checks.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/link-checks.yaml b/.github/workflows/link-checks.yaml index 96cb9370..06ab5b9b 100644 --- a/.github/workflows/link-checks.yaml +++ b/.github/workflows/link-checks.yaml @@ -20,7 +20,7 @@ jobs: contents: read steps: - name: Checkout repository - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v7.0.1 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false - name: Restore lychee cache From ae0cf6590cde4325225885f4b1d2dc6c6bc21459 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 15 Sep 2026 09:30:10 -0700 Subject: [PATCH 07/20] Accommodate adoption of cuda::stream_ref in RAFT handle --- python/nvforest/nvforest/detail/raft_stream.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/nvforest/nvforest/detail/raft_stream.pyx b/python/nvforest/nvforest/detail/raft_stream.pyx index a90de6c2..f3cc504b 100644 --- a/python/nvforest/nvforest/detail/raft_stream.pyx +++ b/python/nvforest/nvforest/detail/raft_stream.pyx @@ -16,7 +16,7 @@ from pylibraft.common.handle cimport handle_t def _get_stream_from_raft_handle(handle: RaftHandle) -> Stream: cdef handle_t* handle_ = handle.getHandle() - cdef cudaStream_t stream = handle_.get_stream() + cdef cudaStream_t stream = handle_.get_stream().get() return Stream.from_handle(stream) From 981c414f217a865c6b4aff49d92409cf65b07336 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 15 Sep 2026 10:02:17 -0700 Subject: [PATCH 08/20] Fix intersphinx for cuda-python --- docs/source/conf.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/source/conf.py b/docs/source/conf.py index afaf139a..f0aedd0c 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -205,6 +205,10 @@ "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 From 1485285aae336d0826e180263ee2f9c3d05dcaef Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 5 Oct 2026 13:56:52 -0700 Subject: [PATCH 09/20] Address reviewer's comment --- cpp/include/nvforest/detail/infer/gpu.cuh | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/cpp/include/nvforest/detail/infer/gpu.cuh b/cpp/include/nvforest/detail/infer/gpu.cuh index 31cc8eaa..4249cff1 100644 --- a/cpp/include/nvforest/detail/infer/gpu.cuh +++ b/cpp/include/nvforest/detail/infer/gpu.cuh @@ -46,10 +46,9 @@ 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_ref stream) { - cuda::stream_ref legacy_default_stream{cudaStreamLegacy}; - cuda::stream_ref per_thread_default_stream{cudaStreamPerThread}; - if (stream != legacy_default_stream && stream != per_thread_default_stream && - stream.device().get() != device.value()) { + if (stream != cuda::stream_ref{cudaStream_t{cudaStreamDefault}} && + stream != cuda::stream_ref{cudaStreamLegacy} && + stream != cuda::stream_ref{cudaStreamPerThread} && 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.device().get())}; From 0a10a92ff13433e4466c61a2395b266a4db8e584 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Fri, 9 Oct 2026 23:19:26 -0700 Subject: [PATCH 10/20] Update after merge --- cpp/tests/treelite_importer.cpp | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/cpp/tests/treelite_importer.cpp b/cpp/tests/treelite_importer.cpp index 37a28037..664ce810 100644 --- a/cpp/tests/treelite_importer.cpp +++ b/cpp/tests/treelite_importer.cpp @@ -502,12 +502,6 @@ TEST(TreeliteImporter, SinglePrecisionDoubleThresholds) { auto constexpr eps = double{std::numeric_limits::epsilon()}; auto constexpr inf = std::numeric_limits::infinity(); -#ifdef NVFOREST_ENABLE_GPU - auto raft_handle = raft::handle_t{}; - auto handle = nvforest::handle_t{raft_handle}; -#else - auto handle = nvforest::handle_t{}; -#endif for (auto threshold : {1.0 + 1.5 * eps, 1.0 + 0.25 * eps, -1e300, 1e300}) { auto tl_model = make_double_threshold_forest(threshold); auto nearest = static_cast(threshold); @@ -532,7 +526,7 @@ TEST(TreeliteImporter, SinglePrecisionDoubleThresholds) auto nvforest_model = import_from_treelite_model(*tl_model, layout, index_type{}, false); ASSERT_FALSE(nvforest_model.is_double_precision()); auto preds = std::vector(X.size(), 0.0f); - nvforest_model.predict(handle, + nvforest_model.predict(nvforest::cuda_stream{}, preds.data(), X.data(), X.size(), From da6d011904b5401883468a4038d62e9cc5939b6a Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 12 Oct 2026 21:44:06 -0700 Subject: [PATCH 11/20] Fix buffer sizing for other prediction type; add cross-memory test --- cpp/include/nvforest/forest_model.hpp | 25 ++++---- cpp/tests/CMakeLists.txt | 1 + cpp/tests/cross_memory.cu | 86 +++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 10 deletions(-) create mode 100644 cpp/tests/cross_memory.cu diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index dea577e5..c6a5a414 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -44,11 +44,14 @@ struct forest_model { decision_forest_); } - /** The number of outputs per row generated by the model */ - auto num_outputs() + /** The number of outputs per row generated by the model for the given type of inference */ + auto num_outputs(infer_kind inference_kind = infer_kind::default_kind) { - return std::visit([](auto&& concrete_forest) { return concrete_forest.num_outputs(); }, - decision_forest_); + return std::visit( + [inference_kind](auto&& concrete_forest) { + return concrete_forest.num_outputs(inference_kind); + }, + decision_forest_); } /** The number of trees in the model */ @@ -120,7 +123,7 @@ struct forest_model { * 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(). + * This must be of size at least ROWS x num_outputs(predict_type). * @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 @@ -162,10 +165,11 @@ struct forest_model { } auto buffer_out = buffer{}; if (output.memory_type() != memory_type()) { - buffer_out = - buffer{row_count * num_outputs(), memory_type(), device_index(), stream}; + buffer_out = buffer{ + row_count * num_outputs(predict_type), memory_type(), device_index(), stream}; } else { - buffer_out = buffer{output.data(), row_count * num_outputs(), memory_type()}; + buffer_out = + buffer{output.data(), row_count * num_outputs(predict_type), memory_type()}; } concrete_forest.predict( buffer_out, buffer_in, stream, predict_type, specified_chunk_size); @@ -228,8 +232,9 @@ struct forest_model { } else { current_device_id = -1; } - 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}; + auto out_buffer = + buffer{output, num_rows * num_outputs(predict_type), out_mem_type, current_device_id}; + auto in_buffer = buffer{input, num_rows * num_features(), in_mem_type, current_device_id}; predict(stream, out_buffer, in_buffer, predict_type, specified_chunk_size); } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 7046f9ba..49373a90 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 CROSS_MEMORY_TEST cross_memory.cu) ConfigureTest(NAME INVALID_STREAM_TEST invalid_stream.cu) endif() ConfigureTest(NAME FOREST_TRAVERSAL_TEST forest/traversal_forest.cpp) diff --git a/cpp/tests/cross_memory.cu b/cpp/tests/cross_memory.cu new file mode 100644 index 00000000..2e21ef9a --- /dev/null +++ b/cpp/tests/cross_memory.cu @@ -0,0 +1,86 @@ +/* + * 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 + +namespace nvforest { + +TEST(CrossMemory, GpuModelWithCpuInputAndOutput) +{ + 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{2, {0, 0}, {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(); + + 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(); + + auto tree_model = model_builder->CommitModel(); + auto nvforest_model = import_from_treelite_model( + *tree_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu); + auto input = std::vector{-1.0f, 1.0f}; + auto const cases = std::array{ + std::pair{infer_kind::default_kind, std::vector{-2.0f, 2.0f}}, + std::pair{infer_kind::per_tree, std::vector{-1.0f, -1.0f, 1.0f, 1.0f}}, + std::pair{infer_kind::leaf_id, std::vector{1.0f, 1.0f, 2.0f, 2.0f}}, + }; + + for (auto const& [infer_type, expected] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto output = std::vector(expected.size()); + nvforest_model.predict(cuda_stream{}, + output.data(), + input.data(), + input.size(), + device_type::cpu, + device_type::cpu, + infer_type); + EXPECT_EQ(output, expected); + } +} + +} // namespace nvforest From 994427810f78d839df2c0d7f7ed5f0db93bc329f Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 12 Oct 2026 21:55:16 -0700 Subject: [PATCH 12/20] Ensure that inference with handle still works --- python/nvforest/tests/test_compat.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python/nvforest/tests/test_compat.py b/python/nvforest/tests/test_compat.py index 4d068a33..6063e216 100644 --- a/python/nvforest/tests/test_compat.py +++ b/python/nvforest/tests/test_compat.py @@ -5,6 +5,8 @@ Compatibility test to ensure that legacy APIs are working """ +import cupy as cp +import numpy as np import pytest import treelite from cuda.core import Device @@ -59,10 +61,21 @@ def test_deprecated_handle(): stream = device.create_stream() handle = nvforest.Handle(stream=stream) + rng = np.random.default_rng(42) + X = rng.standard_normal(size=(1000, 1), dtype="float64") + with pytest.warns( FutureWarning, match=".*`handle` parameter is deprecated and will be removed in 27.02.*", ): - _ = nvforest.load_from_treelite_model( + fm = nvforest.load_from_treelite_model( treelite_model, device="gpu", handle=handle ) + pred_with_handle = cp.asnumpy(fm.predict(X)) + + stream2 = cp.cuda.Stream(non_blocking=True) + fm2 = nvforest.load_from_treelite_model( + treelite_model, device="gpu", stream=stream2 + ) + pred_with_stream = cp.asnumpy(fm2.predict(X)) + np.testing.assert_almost_equal(pred_with_handle, pred_with_stream) From a902d5f8814fa42aa41eed3430650b0cc301358a Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Mon, 12 Oct 2026 22:11:49 -0700 Subject: [PATCH 13/20] Add stream pool interface for predict() --- cpp/include/nvforest/forest_model.hpp | 168 ++++++++++++++++++++ cpp/tests/CMakeLists.txt | 2 +- cpp/tests/cross_memory.cu | 86 ---------- cpp/tests/inference.cu | 220 ++++++++++++++++++++++++++ 4 files changed, 389 insertions(+), 87 deletions(-) delete mode 100644 cpp/tests/cross_memory.cu create mode 100644 cpp/tests/inference.cu diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index c6a5a414..cb3421e2 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -18,6 +18,7 @@ #endif #include +#include #include #include #include @@ -184,6 +185,113 @@ struct forest_model { decision_forest_); } + /** + * Perform inference on given input, using multiple CUDA streams. + * A stream pool lets you overlap memory transfers with partitioned inference. + * + * @param[in] stream_pool_next A function that returns a CUDA stream + * each time it's called. + * @param[in] stream_pool_size A function that returns the number of + * CUDA streams in the pool. + * @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(std::function const& stream_pool_next, + std::function const& stream_pool_size, + 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, + &stream_pool_next, + &stream_pool_size, + &output, + &input, + &specified_chunk_size](auto&& concrete_forest) { + using model_io_t = 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, stream_pool_next(), 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, stream_pool_size()), + 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 = stream_pool_next(); + 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(predict_type), memory_type()}; + } else { + partition_out = + buffer{output.data() + i * partition_size * num_outputs(predict_type), + rows_in_this_partition * num_outputs(predict_type), + 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(predict_type), + 0, + partition_out.size(), + stream); + } + } + } + } else { + throw type_error("Input type does not match model_type"); + } + }, + decision_forest_); + } + /** * Perform inference on given input * @@ -238,6 +346,66 @@ struct forest_model { predict(stream, out_buffer, in_buffer, predict_type, specified_chunk_size); } + /** + * Perform inference on given input, using multiple CUDA streams. + * A stream pool lets you overlap memory transfers with partitioned inference. + * + * @param[in] stream_pool_next A function that returns a CUDA stream + * each time it's called. + * @param[in] stream_pool_size A function that returns the number of + * CUDA streams in the pool. + * @param[out] output Pointer to the memory location where output should end + * up + * @param[in] input Pointer to the input data + * @param[in] num_rows Number of rows in input + * @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] 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(std::function const& stream_pool_next, + std::function const& stream_pool_size, + io_t* output, + io_t* input, + std::size_t num_rows, + device_type out_mem_type, + device_type in_mem_type, + infer_kind predict_type = infer_kind::default_kind, + std::optional specified_chunk_size = std::nullopt) + { + int current_device_id; + if (out_mem_type == device_type::gpu || in_mem_type == device_type::gpu) { +#ifdef NVFOREST_ENABLE_GPU + detail::cuda_check(cudaGetDevice(¤t_device_id)); +#else + throw detail::gpu_unsupported("Tried to use GPU memory in CPU-only build"); +#endif + } else { + current_device_id = -1; + } + 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(stream_pool_next, + stream_pool_size, + out_buffer, + in_buffer, + predict_type, + specified_chunk_size); + } + private: decision_forest_variant decision_forest_; }; diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 49373a90..3eea830e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -81,7 +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 CROSS_MEMORY_TEST cross_memory.cu) + ConfigureTest(NAME INFERENCE_TEST inference.cu) ConfigureTest(NAME INVALID_STREAM_TEST invalid_stream.cu) endif() ConfigureTest(NAME FOREST_TRAVERSAL_TEST forest/traversal_forest.cpp) diff --git a/cpp/tests/cross_memory.cu b/cpp/tests/cross_memory.cu deleted file mode 100644 index 2e21ef9a..00000000 --- a/cpp/tests/cross_memory.cu +++ /dev/null @@ -1,86 +0,0 @@ -/* - * 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 - -namespace nvforest { - -TEST(CrossMemory, GpuModelWithCpuInputAndOutput) -{ - 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{2, {0, 0}, {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(); - - 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(); - - auto tree_model = model_builder->CommitModel(); - auto nvforest_model = import_from_treelite_model( - *tree_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu); - auto input = std::vector{-1.0f, 1.0f}; - auto const cases = std::array{ - std::pair{infer_kind::default_kind, std::vector{-2.0f, 2.0f}}, - std::pair{infer_kind::per_tree, std::vector{-1.0f, -1.0f, 1.0f, 1.0f}}, - std::pair{infer_kind::leaf_id, std::vector{1.0f, 1.0f, 2.0f, 2.0f}}, - }; - - for (auto const& [infer_type, expected] : cases) { - SCOPED_TRACE(static_cast(infer_type)); - auto output = std::vector(expected.size()); - nvforest_model.predict(cuda_stream{}, - output.data(), - input.data(), - input.size(), - device_type::cpu, - device_type::cpu, - infer_type); - EXPECT_EQ(output, expected); - } -} - -} // namespace nvforest diff --git a/cpp/tests/inference.cu b/cpp/tests/inference.cu new file mode 100644 index 00000000..245c4153 --- /dev/null +++ b/cpp/tests/inference.cu @@ -0,0 +1,220 @@ +/* + * 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 + +#include +#include +#include +#include +#include + +namespace nvforest { + +std::unique_ptr get_tree_model() +{ + 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{2, {0, 0}, {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(); + + 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(); + + return model_builder->CommitModel(); +} + +std::vector repeat(std::vector const& x, std::size_t n_repeat) +{ + std::vector ret(x.size() * n_repeat); + for (std::size_t i = 0; i < n_repeat; ++i) { + for (std::size_t j = 0; j < x.size(); j++) { + ret[i * x.size() + j] = x[j]; + } + } + return ret; +} + +template +class GPUModelInferenceTest : public testing::Test {}; + +using BufferTypes = testing::Types, thrust::device_vector>; + +struct BufferTypeNames { + template + static std::string GetName(int) + { + if constexpr (std::is_same_v>) { + return "CPU"; + } else { + return "GPU"; + } + } +}; + +TYPED_TEST_SUITE(GPUModelInferenceTest, BufferTypes, BufferTypeNames); + +TYPED_TEST(GPUModelInferenceTest, Predict) +{ + constexpr auto mem_type = + std::is_same_v> ? device_type::cpu : device_type::gpu; + auto tree_model = get_tree_model(); + auto nvforest_model = import_from_treelite_model( + *tree_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu); + + // Test with a single stream + { + auto input = TypeParam{-1.0f, 1.0f}; + auto const cases = std::array{ + std::pair{infer_kind::default_kind, TypeParam{-2.0f, 2.0f}}, + std::pair{infer_kind::per_tree, TypeParam{-1.0f, -1.0f, 1.0f, 1.0f}}, + std::pair{infer_kind::leaf_id, TypeParam{1.0f, 1.0f, 2.0f, 2.0f}}, + }; + + cuda::device_ref device{0}; + cuda::stream stream{device}; + for (auto const& [infer_type, expected] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto output = TypeParam(expected.size()); + nvforest_model.predict(stream.get(), + thrust::raw_pointer_cast(output.data()), + thrust::raw_pointer_cast(input.data()), + input.size(), + mem_type, + mem_type, + infer_type); + stream.sync(); + EXPECT_EQ(output, expected); + } + } + + // Test with multiple streams + { + std::size_t n_repeat = 256; + auto input = TypeParam(repeat(std::vector{-1.0f, 1.0f}, n_repeat)); + auto const cases = std::array{ + std::pair{infer_kind::default_kind, + TypeParam(repeat(std::vector{-2.0f, 2.0f}, n_repeat))}, + std::pair{infer_kind::per_tree, + TypeParam(repeat(std::vector{-1.0f, -1.0f, 1.0f, 1.0f}, n_repeat))}, + std::pair{infer_kind::leaf_id, + TypeParam(repeat(std::vector{1.0f, 1.0f, 2.0f, 2.0f}, n_repeat))}, + }; + + cuda::device_ref device{0}; + std::size_t n_streams = 4; + std::size_t current_stream_id = 0; + std::vector stream_pool; + for (std::size_t i = 0; i < n_streams; ++i) { + stream_pool.emplace_back(device); + } + auto stream_pool_next = [&]() { + return nvforest::cuda_stream{stream_pool[current_stream_id++].get()}; + }; + auto stream_pool_size = [&]() { return stream_pool.size(); }; + for (auto const& [infer_type, expected] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto output = TypeParam(expected.size()); + nvforest_model.predict(stream_pool_next, + stream_pool_size, + thrust::raw_pointer_cast(output.data()), + thrust::raw_pointer_cast(input.data()), + input.size(), + mem_type, + mem_type, + infer_type); + for (std::size_t i = 0; i < n_streams; ++i) { + stream_pool[i].sync(); + } + EXPECT_EQ(output, expected); + } + } +} + +// TODO(chyunsu3): Remove once we remove the RAFT dependency in 27.02. +TEST(Inference, GPUModelWithRAFTHandle) +{ + auto tree_model = get_tree_model(); + auto nvforest_model = import_from_treelite_model( + *tree_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu); + + std::size_t n_repeat = 256; + auto input = thrust::device_vector(repeat(std::vector{-1.0f, 1.0f}, n_repeat)); + auto const cases = std::array{ + std::pair{infer_kind::default_kind, + thrust::device_vector(repeat(std::vector{-2.0f, 2.0f}, n_repeat))}, + std::pair{ + infer_kind::per_tree, + thrust::device_vector(repeat(std::vector{-1.0f, -1.0f, 1.0f, 1.0f}, n_repeat))}, + std::pair{ + infer_kind::leaf_id, + thrust::device_vector(repeat(std::vector{1.0f, 1.0f, 2.0f, 2.0f}, n_repeat))}, + }; + + std::size_t n_streams = 4; + raft::handle_t handle{cuda::stream_ref{cudaStreamPerThread}, + std::make_shared(n_streams)}; + for (auto const& [infer_type, expected] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto output = std::vector(expected.size()); + nvforest_model.predict([&] { return cuda_stream{handle.get_next_usable_stream().get()}; }, + [&] { return handle.get_stream_pool_size(); }, + thrust::raw_pointer_cast(output.data()), + thrust::raw_pointer_cast(input.data()), + input.size(), + device_type::gpu, + device_type::gpu, + infer_type); + handle.sync_stream_pool(); + EXPECT_EQ(output, expected); + } +} + +} // namespace nvforest From 2642597be82d0e5205220a17ebfda792c735f696 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 13 Oct 2026 20:01:29 -0700 Subject: [PATCH 14/20] Fix gtest crash --- cpp/tests/inference.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/inference.cu b/cpp/tests/inference.cu index 245c4153..7eedb9fb 100644 --- a/cpp/tests/inference.cu +++ b/cpp/tests/inference.cu @@ -203,7 +203,7 @@ TEST(Inference, GPUModelWithRAFTHandle) std::make_shared(n_streams)}; for (auto const& [infer_type, expected] : cases) { SCOPED_TRACE(static_cast(infer_type)); - auto output = std::vector(expected.size()); + auto output = thrust::device_vector(expected.size()); nvforest_model.predict([&] { return cuda_stream{handle.get_next_usable_stream().get()}; }, [&] { return handle.get_stream_pool_size(); }, thrust::raw_pointer_cast(output.data()), From 5c65e166a1d4ffbb895f12dce83b5f33d6d8d8a7 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 13 Oct 2026 20:10:51 -0700 Subject: [PATCH 15/20] Pass pool size as plain std::size_t --- cpp/include/nvforest/forest_model.hpp | 22 ++++++++++------------ cpp/tests/inference.cu | 5 ++--- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index cb3421e2..8c9c5446 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -189,10 +189,9 @@ struct forest_model { * Perform inference on given input, using multiple CUDA streams. * A stream pool lets you overlap memory transfers with partitioned inference. * - * @param[in] stream_pool_next A function that returns a CUDA stream - * each time it's called. - * @param[in] stream_pool_size A function that returns the number of - * CUDA streams in the pool. + * @param[in] stream_pool_next A callback function that returns a CUDA stream + * from the pool. + * @param[in] stream_pool_size The number of CUDA streams in the pool. * @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 @@ -217,7 +216,7 @@ struct forest_model { */ template void predict(std::function const& stream_pool_next, - std::function const& stream_pool_size, + std::size_t stream_pool_size, buffer& output, buffer const& input, infer_kind predict_type = infer_kind::default_kind, @@ -227,7 +226,7 @@ struct forest_model { [this, predict_type, &stream_pool_next, - &stream_pool_size, + stream_pool_size, &output, &input, &specified_chunk_size](auto&& concrete_forest) { @@ -242,7 +241,7 @@ struct forest_model { auto row_count = input.size() / num_features(); auto partition_size = - std::max(detail::ceildiv(row_count, stream_pool_size()), + std::max(detail::ceildiv(row_count, stream_pool_size), 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) { @@ -350,10 +349,9 @@ struct forest_model { * Perform inference on given input, using multiple CUDA streams. * A stream pool lets you overlap memory transfers with partitioned inference. * - * @param[in] stream_pool_next A function that returns a CUDA stream - * each time it's called. - * @param[in] stream_pool_size A function that returns the number of - * CUDA streams in the pool. + * @param[in] stream_pool_next A callback function that returns a CUDA stream + * from the pool. + * @param[in] stream_pool_size The number of CUDA streams in the pool. * @param[out] output Pointer to the memory location where output should end * up * @param[in] input Pointer to the input data @@ -377,7 +375,7 @@ struct forest_model { */ template void predict(std::function const& stream_pool_next, - std::function const& stream_pool_size, + std::size_t const& stream_pool_size, io_t* output, io_t* input, std::size_t num_rows, diff --git a/cpp/tests/inference.cu b/cpp/tests/inference.cu index 7eedb9fb..b3c7014b 100644 --- a/cpp/tests/inference.cu +++ b/cpp/tests/inference.cu @@ -158,12 +158,11 @@ TYPED_TEST(GPUModelInferenceTest, Predict) auto stream_pool_next = [&]() { return nvforest::cuda_stream{stream_pool[current_stream_id++].get()}; }; - auto stream_pool_size = [&]() { return stream_pool.size(); }; for (auto const& [infer_type, expected] : cases) { SCOPED_TRACE(static_cast(infer_type)); auto output = TypeParam(expected.size()); nvforest_model.predict(stream_pool_next, - stream_pool_size, + n_streams, thrust::raw_pointer_cast(output.data()), thrust::raw_pointer_cast(input.data()), input.size(), @@ -205,7 +204,7 @@ TEST(Inference, GPUModelWithRAFTHandle) SCOPED_TRACE(static_cast(infer_type)); auto output = thrust::device_vector(expected.size()); nvforest_model.predict([&] { return cuda_stream{handle.get_next_usable_stream().get()}; }, - [&] { return handle.get_stream_pool_size(); }, + handle.get_stream_pool_size(), thrust::raw_pointer_cast(output.data()), thrust::raw_pointer_cast(input.data()), input.size(), From 20b8dd9ec6487b1d53ed65e2e7232154f3dd4be9 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 13 Oct 2026 20:15:58 -0700 Subject: [PATCH 16/20] Document that pool size must be >0 --- 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 8c9c5446..1c2b7d5d 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -191,7 +191,8 @@ struct forest_model { * * @param[in] stream_pool_next A callback function that returns a CUDA stream * from the pool. - * @param[in] stream_pool_size The number of CUDA streams in the pool. + * @param[in] stream_pool_size The number of CUDA streams in the pool. This + * must be greater than zero. * @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 @@ -351,7 +352,8 @@ struct forest_model { * * @param[in] stream_pool_next A callback function that returns a CUDA stream * from the pool. - * @param[in] stream_pool_size The number of CUDA streams in the pool. + * @param[in] stream_pool_size The number of CUDA streams in the pool. This + * must be greater than zero. * @param[out] output Pointer to the memory location where output should end * up * @param[in] input Pointer to the input data From 72317e0cdc16e4c2ea955ba0f54cd975b2771cbb Mon Sep 17 00:00:00 2001 From: Simon Adorf Date: Fri, 16 Oct 2026 13:45:36 -0500 Subject: [PATCH 17/20] Validate stream pool size and pass it by value --- cpp/include/nvforest/forest_model.hpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index 1c2b7d5d..e2e380c3 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -196,7 +197,7 @@ struct forest_model { * @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(). + * to this output location. This must be of size at least ROWS x num_outputs(predict_type). * @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 @@ -223,6 +224,9 @@ struct forest_model { infer_kind predict_type = infer_kind::default_kind, std::optional specified_chunk_size = std::nullopt) { + if (stream_pool_size == 0) { + throw std::invalid_argument("stream_pool_size must be greater than zero"); + } std::visit( [this, predict_type, @@ -377,7 +381,7 @@ struct forest_model { */ template void predict(std::function const& stream_pool_next, - std::size_t const& stream_pool_size, + std::size_t stream_pool_size, io_t* output, io_t* input, std::size_t num_rows, From fc362793156d0ec19a624c52eb3a4ce456766408 Mon Sep 17 00:00:00 2001 From: Simon Adorf Date: Fri, 16 Oct 2026 13:45:56 -0500 Subject: [PATCH 18/20] Preserve inference buffer devices, streams, and output widths --- cpp/include/nvforest/forest_model.hpp | 29 ++++++++++++++++++--------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/cpp/include/nvforest/forest_model.hpp b/cpp/include/nvforest/forest_model.hpp index e2e380c3..346a48e4 100644 --- a/cpp/include/nvforest/forest_model.hpp +++ b/cpp/include/nvforest/forest_model.hpp @@ -163,15 +163,18 @@ struct forest_model { 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()}; + buffer_in = buffer{ + input.data(), row_count * num_features(), memory_type(), input.device_index()}; } auto buffer_out = buffer{}; if (output.memory_type() != memory_type()) { buffer_out = buffer{ row_count * num_outputs(predict_type), memory_type(), device_index(), stream}; } else { - buffer_out = - buffer{output.data(), row_count * num_outputs(predict_type), memory_type()}; + buffer_out = buffer{output.data(), + row_count * num_outputs(predict_type), + memory_type(), + output.device_index()}; } concrete_forest.predict( buffer_out, buffer_in, stream, predict_type, specified_chunk_size); @@ -255,7 +258,8 @@ struct forest_model { 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()}; + partition_in = buffer{ + rows_in_this_partition * num_features(), memory_type(), device_index(), stream}; copy(partition_in, input, 0, @@ -265,17 +269,21 @@ struct forest_model { } else { partition_in = buffer{input.data() + i * partition_size * num_features(), rows_in_this_partition * num_features(), - memory_type()}; + memory_type(), + input.device_index()}; } auto partition_out = buffer{}; if (output.memory_type() != memory_type()) { - partition_out = - buffer{rows_in_this_partition * num_outputs(predict_type), memory_type()}; + partition_out = buffer{rows_in_this_partition * num_outputs(predict_type), + memory_type(), + device_index(), + stream}; } else { partition_out = buffer{output.data() + i * partition_size * num_outputs(predict_type), rows_in_this_partition * num_outputs(predict_type), - memory_type()}; + memory_type(), + output.device_index()}; } concrete_forest.predict( partition_out, partition_in, stream, predict_type, specified_chunk_size); @@ -400,8 +408,9 @@ struct forest_model { } else { current_device_id = -1; } - 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}; + auto out_buffer = + buffer{output, num_rows * num_outputs(predict_type), out_mem_type, current_device_id}; + auto in_buffer = buffer{input, num_rows * num_features(), in_mem_type, current_device_id}; predict(stream_pool_next, stream_pool_size, out_buffer, From 922617f678e50ab718b67b1a4e727d9a077709ad Mon Sep 17 00:00:00 2001 From: Simon Adorf Date: Fri, 16 Oct 2026 13:52:17 -0500 Subject: [PATCH 19/20] Exercise partitioned inference across memory locations and devices --- cpp/tests/inference.cu | 198 +++++++++++++++++++++++------------------ 1 file changed, 109 insertions(+), 89 deletions(-) diff --git a/cpp/tests/inference.cu b/cpp/tests/inference.cu index b3c7014b..a849c960 100644 --- a/cpp/tests/inference.cu +++ b/cpp/tests/inference.cu @@ -4,6 +4,7 @@ */ #include +#include #include #include #include @@ -21,8 +22,8 @@ #include #include -#include -#include +#include +#include #include #include @@ -82,101 +83,120 @@ std::vector repeat(std::vector const& x, std::size_t n_repeat) return ret; } -template -class GPUModelInferenceTest : public testing::Test {}; +// Exercise every input/output memory combination on GPU 0 and, when available, GPU 1. +class GPUModelInferenceTest + : public testing::TestWithParam> {}; -using BufferTypes = testing::Types, thrust::device_vector>; - -struct BufferTypeNames { - template - static std::string GetName(int) - { - if constexpr (std::is_same_v>) { - return "CPU"; - } else { - return "GPU"; - } - } -}; - -TYPED_TEST_SUITE(GPUModelInferenceTest, BufferTypes, BufferTypeNames); - -TYPED_TEST(GPUModelInferenceTest, Predict) +TEST_P(GPUModelInferenceTest, Predict) { - constexpr auto mem_type = - std::is_same_v> ? device_type::cpu : device_type::gpu; - auto tree_model = get_tree_model(); - auto nvforest_model = import_from_treelite_model( - *tree_model, tree_layout::breadth_first, index_type{}, false, device_type::gpu); - - // Test with a single stream - { - auto input = TypeParam{-1.0f, 1.0f}; - auto const cases = std::array{ - std::pair{infer_kind::default_kind, TypeParam{-2.0f, 2.0f}}, - std::pair{infer_kind::per_tree, TypeParam{-1.0f, -1.0f, 1.0f, 1.0f}}, - std::pair{infer_kind::leaf_id, TypeParam{1.0f, 1.0f, 2.0f, 2.0f}}, - }; - - cuda::device_ref device{0}; - cuda::stream stream{device}; - for (auto const& [infer_type, expected] : cases) { - SCOPED_TRACE(static_cast(infer_type)); - auto output = TypeParam(expected.size()); - nvforest_model.predict(stream.get(), - thrust::raw_pointer_cast(output.data()), - thrust::raw_pointer_cast(input.data()), - input.size(), - mem_type, - mem_type, - infer_type); - stream.sync(); - EXPECT_EQ(output, expected); - } + auto const [device_id, in_mem_type, out_mem_type] = GetParam(); + int device_count; + detail::cuda_check(cudaGetDeviceCount(&device_count)); + if (device_id >= device_count) { GTEST_SKIP() << "Requires GPU " << device_id; } + auto device_context = detail::device_setter{detail::device_id{device_id}}; + + constexpr std::size_t n_streams = 4; + std::vector stream_pool; + for (std::size_t i = 0; i < n_streams; ++i) { + stream_pool.emplace_back(cuda::device_ref{device_id}); } - - // Test with multiple streams - { - std::size_t n_repeat = 256; - auto input = TypeParam(repeat(std::vector{-1.0f, 1.0f}, n_repeat)); - auto const cases = std::array{ - std::pair{infer_kind::default_kind, - TypeParam(repeat(std::vector{-2.0f, 2.0f}, n_repeat))}, - std::pair{infer_kind::per_tree, - TypeParam(repeat(std::vector{-1.0f, -1.0f, 1.0f, 1.0f}, n_repeat))}, - std::pair{infer_kind::leaf_id, - TypeParam(repeat(std::vector{1.0f, 1.0f, 2.0f, 2.0f}, n_repeat))}, - }; - - cuda::device_ref device{0}; - std::size_t n_streams = 4; - std::size_t current_stream_id = 0; - std::vector stream_pool; - for (std::size_t i = 0; i < n_streams; ++i) { - stream_pool.emplace_back(device); - } - auto stream_pool_next = [&]() { - return nvforest::cuda_stream{stream_pool[current_stream_id++].get()}; - }; - for (auto const& [infer_type, expected] : cases) { - SCOPED_TRACE(static_cast(infer_type)); - auto output = TypeParam(expected.size()); - nvforest_model.predict(stream_pool_next, - n_streams, - thrust::raw_pointer_cast(output.data()), - thrust::raw_pointer_cast(input.data()), - input.size(), - mem_type, - mem_type, - infer_type); - for (std::size_t i = 0; i < n_streams; ++i) { - stream_pool[i].sync(); + auto tree_model = get_tree_model(); + auto nvforest_model = import_from_treelite_model(*tree_model, + tree_layout::breadth_first, + index_type{}, + false, + device_type::gpu, + device_id, + stream_pool[0].get()); + // Model initialization must finish before inference can use the other streams. + stream_pool[0].sync(); + + // Larger than four minimum-size partitions, with an uneven final partition. + constexpr std::size_t n_repeat = 16385; + auto host_input = repeat(std::vector{-1.0f, 1.0f}, n_repeat); + auto input_view = buffer{host_input.data(), host_input.size()}; + auto input = buffer{input_view, in_mem_type, device_id, stream_pool[0].get()}; + stream_pool[0].sync(); + auto const cases = std::array{ + std::pair{infer_kind::default_kind, std::vector{-2.0f, 2.0f}}, + std::pair{infer_kind::per_tree, std::vector{-1.0f, -1.0f, 1.0f, 1.0f}}, + std::pair{infer_kind::leaf_id, std::vector{1.0f, 1.0f, 2.0f, 2.0f}}, + }; + std::size_t stream_calls = 0; + auto stream_pool_next = [&]() { + return cuda_stream{stream_pool[stream_calls++ % n_streams].get()}; + }; + for (auto const& [infer_type, row_pair] : cases) { + SCOPED_TRACE(static_cast(infer_type)); + auto expected = repeat(row_pair, n_repeat); + for (bool use_pool : {false, true}) { + SCOPED_TRACE(use_pool); + for (bool use_buffers : {false, true}) { + SCOPED_TRACE(use_buffers); + auto output = buffer{expected.size(), out_mem_type, device_id, stream_pool[0].get()}; + stream_pool[0].sync(); + auto const previous_stream_calls = stream_calls; + if (use_pool) { + if (use_buffers) { + nvforest_model.predict(stream_pool_next, n_streams, output, input, infer_type); + } else { + nvforest_model.predict(stream_pool_next, + n_streams, + output.data(), + input.data(), + host_input.size(), + out_mem_type, + in_mem_type, + infer_type); + } + auto const cross_memory = + in_mem_type != device_type::gpu || out_mem_type != device_type::gpu; + EXPECT_EQ(stream_calls - previous_stream_calls, cross_memory ? n_streams : 1); + } else if (use_buffers) { + nvforest_model.predict(stream_pool[0].get(), output, input, infer_type); + } else { + nvforest_model.predict(stream_pool[0].get(), + output.data(), + input.data(), + host_input.size(), + out_mem_type, + in_mem_type, + infer_type); + } + for (auto& stream : stream_pool) { + stream.sync(); + } + auto actual = buffer{output, device_type::cpu, 0, stream_pool[0].get()}; + stream_pool[0].sync(); + EXPECT_EQ(std::vector(actual.data(), actual.data() + actual.size()), expected); } - EXPECT_EQ(output, expected); } } } +INSTANTIATE_TEST_SUITE_P(MemoryLocations, + GPUModelInferenceTest, + testing::Combine(testing::Values(0, 1), + testing::Values(device_type::cpu, device_type::gpu), + testing::Values(device_type::cpu, device_type::gpu))); + +TEST(Inference, EmptyStreamPool) +{ + auto tree_model = get_tree_model(); + auto model = import_from_treelite_model(*tree_model, tree_layout::breadth_first); + auto input = buffer{1}; + auto output = buffer{1}; + auto next_stream = []() -> cuda_stream { + ADD_FAILURE() << "An empty stream pool must be rejected before requesting a stream"; + return cuda_stream{}; + }; + EXPECT_THROW(model.predict(next_stream, 0, output, input), std::invalid_argument); + EXPECT_THROW( + model.predict( + next_stream, 0, output.data(), input.data(), 1, device_type::cpu, device_type::cpu), + std::invalid_argument); +} + // TODO(chyunsu3): Remove once we remove the RAFT dependency in 27.02. TEST(Inference, GPUModelWithRAFTHandle) { From 3e7aafb9fa04efc80c272dee75238f60fc910324 Mon Sep 17 00:00:00 2001 From: Simon Adorf Date: Fri, 16 Oct 2026 15:39:46 -0500 Subject: [PATCH 20/20] Import model on a valid stream in wrong-device prediction test --- cpp/tests/invalid_stream.cu | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/cpp/tests/invalid_stream.cu b/cpp/tests/invalid_stream.cu index 0d9b028b..2777472e 100644 --- a/cpp/tests/invalid_stream.cu +++ b/cpp/tests/invalid_stream.cu @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -53,21 +54,25 @@ TEST(InvalidStream, wrong_device) model_builder->EndNode(); model_builder->EndTree(); - cuda_stream stream; - { - auto device_context = detail::device_setter{detail::device_id{1}}; - detail::cuda_check(cudaStreamCreate(&stream)); - } + auto device_context = detail::device_setter{detail::device_id{0}}; + cuda::stream import_stream{cuda::device_ref{0}}; + cuda::stream wrong_device_stream{cuda::device_ref{1}}; - // Loading a model on device 0 and inferencing on device 1 is an error. + // Import on a valid stream before testing prediction with a wrong-device stream. 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 nvforest_model = import_from_treelite_model(*tl_model, + tree_layout::breadth_first, + index_type{}, + false, + device_type::gpu, + 0, + import_stream.get()); + import_stream.sync(); auto input = thrust::device_vector{1.0f}; auto output = thrust::device_vector(1); EXPECT_THAT( [&]() { - nvforest_model.predict(stream, + nvforest_model.predict(wrong_device_stream.get(), thrust::raw_pointer_cast(output.data()), thrust::raw_pointer_cast(input.data()), 1,