Remove RMM dependency and prepare for RAFT removal, Part 2

conda/recipes/nvforest/recipe.yaml

@@ -99,6 +99,7 @@ requirements:
     - numpy >=2.0,<3.0a0
     - scikit-learn >=1.4
     - pylibraft =${{ minor_version }}
+    - cuda-core >=1.0.0,<2.0.0
     - python
     - treelite ${{ treelite_version_run }}
     - cuda-cudart

cpp/include/nvforest/Implementation.md

@@ -147,13 +147,16 @@ a standard benchmark) on the CPU.
 
 With some motivation for the general approach to CPU-GPU interoperability, we
 now offer an overview of the layout of the codebase to help guide future
-improvements. Because `nvforest::detail` utilities are going to be moved to RAFT or other
-general-purpose libraries, we will not review anything within the `nvforest::detail`
-directory here.
+improvements. The `nvforest::detail` namespace is private implementation infrastructure and
+may change without notice, so this overview focuses on its responsibilities
+rather than treating individual utilities as stable APIs.
 
 ### Public Headers
 * `constants.hpp`: Contains constant values that may be useful in working
   with nvForest in other C++ applications
+* `cuda_stream.hpp`: Defines the public `nvforest::cuda_stream` abstraction.
+  It aliases `cudaStream_t` in GPU builds and is a CUDA-free placeholder in
+  CPU-only builds.
 * `decision_forest.hpp`: Provides `decision_forest`, a template which provides
   concrete implementations of a decision forest. Because different types may
   be optimal for different sizes of models or models with different features,

docs/source/build.rst

@@ -349,10 +349,6 @@ nvForest's cmake has the following configurable flags available:
      - [ON, OFF]
      - OFF
      - Whether to cache build artifacts with ccache.
-   * - NVFOREST_USE_RAFT_STATIC
-     - [ON, OFF]
-     - OFF
-     - Whether to statically link the RAFT library.
    * - NVFOREST_USE_TREELITE_STATIC
      - [ON, OFF]
      - OFF

docs/source/cpp_api.rst

@@ -37,6 +37,8 @@ Enums and constants
 Type aliases
 ============
 
+.. doxygentypedef:: nvforest::cuda_stream
+
 .. doxygentypedef:: nvforest::decision_forest_variant
 
 .. doxygentypedef:: nvforest::detail::preset_decision_forest

docs/source/python_api.rst

@@ -2,6 +2,9 @@
 Python API
 ##########
 
+Model loading functions
+=======================
+
 .. role:: py(code)
    :language: python
    :class: highlight

python/nvforest/nvforest/_base.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
@@ -34,9 +34,7 @@ def predict_per_tree(
             to the correct datatype before inference. If this input is in a
             memory location that is inaccessible to the current device type
             (as set with the 'device' parameter in the constructor),
-            it will be copied to the correct location. This copy will be
-            distributed across as many CUDA streams as are available
-            in the stream pool of the model's RAFT handle.
+            it will be copied to the correct location on the model's CUDA stream.
         chunk_size :
             The number of rows to simultaneously process in one iteration
             of the inference algorithm. Batches are further broken down into
@@ -72,9 +70,7 @@ def apply(
             to the correct datatype before inference. If this input is in a
             memory location that is inaccessible to the current device type
             (as set with the 'device' parameter in the constructor),
-            it will be copied to the correct location. This copy will be
-            distributed across as many CUDA streams as are available
-            in the stream pool of the model's RAFT handle.
+            it will be copied to the correct location on the model's CUDA stream.
         chunk_size :
             The number of rows to simultaneously process in one iteration
             of the inference algorithm. Batches are further broken down into
@@ -226,9 +222,7 @@ def predict(
             to the correct datatype before inference. If this input is in a
             memory location that is inaccessible to the current device type
             (as set with the 'device' parameter in the constructor),
-            it will be copied to the correct location. This copy will be
-            distributed across as many CUDA streams as are available
-            in the stream pool of the model's RAFT handle.
+            it will be copied to the correct location on the model's CUDA stream.
         chunk_size :
             The number of rows to simultaneously process in one iteration
             of the inference algorithm. Batches are further broken down into
@@ -270,9 +264,7 @@ def predict_proba(
             to the correct datatype before inference. If this input is in a
             memory location that is inaccessible to the current device type
             (as set with the 'device' parameter in the constructor),
-            it will be copied to the correct location. This copy will be
-            distributed across as many CUDA streams as are available
-            in the stream pool of the model's RAFT handle.
+            it will be copied to the correct location on the model's CUDA stream.
         chunk_size :
             The number of rows to simultaneously process in one iteration
             of the inference algorithm. Batches are further broken down into
@@ -314,9 +306,7 @@ def predict(
             to the correct datatype before inference. If this input is in a
             memory location that is inaccessible to the current device type
             (as set with the 'device' parameter in the constructor),
-            it will be copied to the correct location. This copy will be
-            distributed across as many CUDA streams as are available
-            in the stream pool of the model's RAFT handle.
+            it will be copied to the correct location on the model's CUDA stream.
         chunk_size :
             The number of rows to simultaneously process in one iteration
             of the inference algorithm. Batches are further broken down into

python/nvforest/tests/test_nvforest.py

@@ -10,6 +10,7 @@
 import pandas as pd
 import pytest
 import treelite
+from cuda.core import Device, Stream
 
 # Import XGBoost before scikit-learn to work around a libgomp bug
 # See https://github.com/dmlc/xgboost/issues/7110
@@ -926,22 +927,13 @@ def test_missing_categorical(category_list):
     input = np.array([[np.nan]])
     gtil_preds = treelite.gtil.predict(model, input)
     fm = nvforest.load_from_treelite_model(model)
-    fil_preds = _get_numpy_array(fm.predict(input))
-    np.testing.assert_equal(fil_preds.flatten(), gtil_preds.flatten())
+    nvforest_preds = _get_numpy_array(fm.predict(input))
+    np.testing.assert_equal(nvforest_preds.flatten(), gtil_preds.flatten())
 
 
 @pytest.mark.parametrize("device_id", [None, 0, 1, 2])
 @pytest.mark.parametrize("model_kind", ["sklearn", "xgboost"])
 def test_device_selection(device_id, model_kind, tmp_path):
-    import sklearn
-    from packaging.version import Version
-
-    # TODO(hcho3): Remove this once Rapids adopts XGBoost 3.1.3
-    if model_kind == "xgboost" and Version(sklearn.__version__) >= Version(
-        "1.8.0.dev0"
-    ):
-        pytest.skip("xgboost is incompatible with sklearn >= 1.8.0.dev0")
-
     current_device = cp.cuda.runtime.getDevice()
 
     if device_id is not None and device_id >= cp.cuda.runtime.getDeviceCount():
@@ -1018,6 +1010,36 @@ def test_device_selection(device_id, model_kind, tmp_path):
     assert cp.cuda.runtime.getDevice() == current_device
 
 
+def test_cuda_core_stream():
+    previous_device = Device()
+    try:
+        device = Device(0)
+        device.set_current()
+        stream = device.create_stream()
+
+        X, y = _simulate_data(
+            100, 8, n_informative=2, random_state=0, classification=True
+        )
+        skl_model = RandomForestClassifier(
+            max_depth=3, random_state=0, n_estimators=5
+        ).fit(X, y)
+        default_fm = nvforest.load_from_sklearn(
+            skl_model, device="gpu", device_id=0
+        )
+        fm = nvforest.load_from_sklearn(
+            skl_model, device="gpu", device_id=0, stream=stream
+        )
+
+        assert isinstance(default_fm.forest.stream, Stream)
+        assert isinstance(stream, Stream)
+        assert fm.forest.stream is stream
+        np.testing.assert_array_equal(
+            cp.asnumpy(fm.predict(X)), skl_model.predict(X)
+        )
+    finally:
+        previous_device.set_current()
+
+
 def test_wide_data():
     n_rows = 50
     n_features = 100000