Migrate RMM usage to CCCL MR design

ci/build_cpp.sh

@@ -1,10 +1,11 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 source rapids-configure-sccache
 

ci/build_docs.sh

@@ -1,9 +1,10 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 rapids-logger "Create test conda environment"
 . /opt/conda/etc/profile.d/conda.sh
@@ -23,6 +24,7 @@ PYTHON_CHANNEL=$(rapids-download-conda-from-github python)
 rapids-logger "Generating conda environment YAML"
 
 rapids-dependency-file-generator \
+  "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \
   --output conda \
   --file-key docs \
   --prepend-channel "${CPP_CHANNEL}" \

ci/build_python.sh

@@ -1,9 +1,10 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 source rapids-configure-sccache
 

ci/build_wheel_cuopt.sh

@@ -5,9 +5,10 @@
 
 
 set -euo pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 # Install rockylinux repo
 if command -v dnf &> /dev/null; then
     bash ci/utils/update_rockylinux_repo.sh

ci/build_wheel_cuopt_mps_parser.sh

@@ -5,9 +5,10 @@
 
 
 set -euo pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 package_dir="python/cuopt/cuopt/linear_programming"
 
 export SKBUILD_CMAKE_ARGS="-DCUOPT_BUILD_WHEELS=ON;-DDISABLE_DEPRECATION_WARNINGS=ON"

ci/build_wheel_cuopt_server.sh

@@ -4,9 +4,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 package_dir="python/cuopt_server"
 
 ci/build_wheel.sh cuopt_server ${package_dir}

ci/build_wheel_cuopt_sh_client.sh

@@ -4,9 +4,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 package_dir="python/cuopt_self_hosted"
 
 ci/build_wheel.sh cuopt_sh_client ${package_dir}

ci/build_wheel_libcuopt.sh

@@ -3,9 +3,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 package_name="libcuopt"
 package_dir="python/libcuopt"
 

ci/test_cpp.sh

@@ -1,9 +1,10 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 . /opt/conda/etc/profile.d/conda.sh
 
@@ -14,6 +15,7 @@ CPP_CHANNEL=$(rapids-download-conda-from-github cpp)
 
 rapids-logger "Generate C++ testing dependencies"
 rapids-dependency-file-generator \
+  "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \
   --output conda \
   --file-key test_cpp \
   --prepend-channel "${CPP_CHANNEL}" \

ci/test_cpp_memcheck.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 if [[ "$(date +%A)" != "Friday" ]]; then
@@ -9,6 +9,7 @@ if [[ "$(date +%A)" != "Friday" ]]; then
 fi
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 . /opt/conda/etc/profile.d/conda.sh
 
@@ -19,6 +20,7 @@ RAPIDS_VERSION="$(rapids-version)"
 
 rapids-logger "Generate C++ testing dependencies"
 rapids-dependency-file-generator \
+  "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \
   --output conda \
   --file-key test_cpp \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" | tee env.yaml

ci/test_notebooks.sh

@@ -1,9 +1,10 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 . /opt/conda/etc/profile.d/conda.sh
 
@@ -19,6 +20,7 @@ rapids-logger "Generate notebook testing dependencies"
 ENV_YAML_DIR="$(mktemp -d)"
 
 rapids-dependency-file-generator \
+  "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \
   --output conda \
   --file-key test_notebooks \
   --prepend-channel "${CPP_CHANNEL}" \

ci/test_python.sh

@@ -4,6 +4,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 . /opt/conda/etc/profile.d/conda.sh
 
@@ -16,6 +17,7 @@ PYTHON_CHANNEL=$(rapids-download-conda-from-github python)
 
 rapids-logger "Generate Python testing dependencies"
 rapids-dependency-file-generator \
+  "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \
   --output conda \
   --file-key test_python \
   --prepend-channel "${CPP_CHANNEL}" \

ci/test_self_hosted_service.sh

@@ -4,9 +4,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 # Download the cuopt built in the previous step
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
 CUOPT_MPS_PARSER_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cuopt_mps_parser" rapids-download-wheels-from-github python)

ci/test_skills_assets.sh

@@ -9,6 +9,7 @@
 # CLI: run cuopt_cli on each sample .mps in API-CLI skill assets.
 
 set -euo pipefail
+source ./ci/use_conda_packages_from_prs.sh
 
 # Use rapids-logger in CI; fall back to echo for local testing
 if command -v rapids-logger &>/dev/null; then

ci/test_wheel_cuopt.sh

@@ -8,6 +8,7 @@ set -euo pipefail
 # sets up a constraints file for 'pip' and puts its location in an exported variable PIP_EXPORT,
 # so those constraints will affect all future 'pip install' calls
 source rapids-init-pip
+source ./ci/use_wheels_from_prs.sh
 
 # Download the packages built in the previous step
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"

ci/test_wheel_cuopt_server.sh

@@ -4,9 +4,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 set -eou pipefail
-
 source rapids-init-pip
 
+source ./ci/use_wheels_from_prs.sh
+
 # Download the packages built in the previous step
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
 CUOPT_MPS_PARSER_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cuopt_mps_parser" rapids-download-wheels-from-github python)

ci/use_conda_packages_from_prs.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+# download CI artifacts
+LIBRMM_CHANNEL=$(rapids-get-pr-artifact rmm 2361 cpp conda)
+RMM_CHANNEL=$(rapids-get-pr-artifact rmm 2361 python conda --stable)
+LIBUCXX_CHANNEL=$(rapids-get-pr-artifact ucxx 636 cpp conda)
+UCXX_CHANNEL=$(rapids-get-pr-artifact ucxx 636 python conda --stable)
+LIBRAFT_CHANNEL=$(rapids-get-pr-artifact raft 2996 cpp conda)
+RAFT_CHANNEL=$(rapids-get-pr-artifact raft 2996 python conda --stable)
+
+RAPIDS_PREPENDED_CONDA_CHANNELS=(
+    "${LIBRMM_CHANNEL}"
+    "${RMM_CHANNEL}"
+    "${LIBUCXX_CHANNEL}"
+    "${UCXX_CHANNEL}"
+    "${LIBRAFT_CHANNEL}"
+    "${RAFT_CHANNEL}"
+)
+export RAPIDS_PREPENDED_CONDA_CHANNELS
+
+RAPIDS_EXTRA_CONDA_CHANNEL_ARGS=()
+for _channel in "${RAPIDS_PREPENDED_CONDA_CHANNELS[@]}"
+do
+   conda config --system --add channels "${_channel}"
+   RAPIDS_EXTRA_CONDA_CHANNEL_ARGS+=(--prepend-channel "${_channel}")
+done

ci/use_wheels_from_prs.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+RAPIDS_PY_CUDA_SUFFIX=$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")
+
+LIBRMM_WHEELHOUSE=$(
+  RAPIDS_PY_WHEEL_NAME="librmm_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-artifact rmm 2361 cpp wheel
+)
+RMM_WHEELHOUSE=$(
+  rapids-get-pr-artifact rmm 2361 python wheel --stable
+)
+LIBUCXX_WHEELHOUSE=$(
+  RAPIDS_PY_WHEEL_NAME="libucxx_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-artifact ucxx 636 cpp wheel
+)
+UCXX_WHEELHOUSE=$(
+  rapids-get-pr-artifact ucxx 636 python wheel --stable
+)
+LIBRAFT_WHEELHOUSE=$(
+  RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-artifact raft 2996 cpp wheel
+)
+RAFT_WHEELHOUSE=$(
+  rapids-get-pr-artifact raft 2996 python wheel --stable --pkg_name pylibraft
+)
+
+cat >> "${PIP_CONSTRAINT}" <<EOF
+librmm-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBRMM_WHEELHOUSE}"/librmm_*.whl)
+rmm-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${RMM_WHEELHOUSE}"/rmm_*.whl)
+libucxx-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBUCXX_WHEELHOUSE}"/libucxx_*.whl)
+ucxx-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${UCXX_WHEELHOUSE}"/ucxx_*.whl)
+libraft-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBRAFT_WHEELHOUSE}"/libraft_*.whl)
+pylibraft-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${RAFT_WHEELHOUSE}"/pylibraft_*.whl)
+EOF

cpp/cuopt_cli.cpp

@@ -415,15 +415,16 @@ int main(int argc, char* argv[])
 
   // Only initialize CUDA resources if using GPU memory backend (not remote execution)
   auto memory_backend = cuopt::linear_programming::get_memory_backend_type();
-  std::vector<std::shared_ptr<rmm::mr::device_memory_resource>> memory_resources;
+  std::vector<rmm::mr::cuda_async_memory_resource> memory_resources;
 
   if (memory_backend == cuopt::linear_programming::memory_backend_t::GPU) {
     const int num_gpus = settings.get_parameter<int>(CUOPT_NUM_GPUS);
 
+    memory_resources.reserve(std::min(raft::device_setter::get_device_count(), num_gpus));
     for (int i = 0; i < std::min(raft::device_setter::get_device_count(), num_gpus); ++i) {
       RAFT_CUDA_TRY(cudaSetDevice(i));
-      memory_resources.push_back(make_async());
-      rmm::mr::set_per_device_resource(rmm::cuda_device_id{i}, memory_resources.back().get());
+      memory_resources.emplace_back();
+      rmm::mr::set_per_device_resource(rmm::cuda_device_id{i}, memory_resources.back());
     }
     RAFT_CUDA_TRY(cudaSetDevice(0));
   }

cpp/src/barrier/sparse_cholesky.cuh

@@ -247,8 +247,8 @@ class sparse_cholesky_cudss_t : public sparse_cholesky_base_t<i_t, f_t> {
     CUDSS_CALL_AND_CHECK_EXIT(cudssSetStream(handle, stream), status, "cudaStreamCreate");
 
     mem_handler.ctx          = reinterpret_cast<void*>(handle_ptr_->get_workspace_resource());
-    mem_handler.device_alloc = cudss_device_alloc<rmm::mr::device_memory_resource>;
-    mem_handler.device_free  = cudss_device_dealloc<rmm::mr::device_memory_resource>;
+    mem_handler.device_alloc = cudss_device_alloc<void>;
+    mem_handler.device_free  = cudss_device_dealloc<void>;
 
     CUDSS_CALL_AND_CHECK_EXIT(
       cudssSetDeviceMemHandler(handle, &mem_handler), status, "cudssSetDeviceMemHandler");

cpp/src/pdlp/termination_strategy/infeasibility_information.cu

@@ -15,6 +15,8 @@
 
 #include <mip_heuristics/mip_constants.hpp>
 
+#include <thrust/iterator/transform_output_iterator.h>
+
 #include <raft/sparse/detail/cusparse_wrappers.h>
 #include <raft/core/nvtx.hpp>
 #include <raft/linalg/detail/cublas_wrappers.hpp>

cpp/src/routing/ges_solver.cu

@@ -16,8 +16,6 @@
 #include "adapters/assignment_adapter.cuh"
 #include "ges/guided_ejection_search.cuh"
 
-#include <rmm/mr/device_memory_resource.hpp>
-
 namespace cuopt {
 namespace routing {
 

cpp/src/utilities/cuda_helpers.cuh

@@ -16,8 +16,6 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/mr/cuda_async_memory_resource.hpp>
-#include <rmm/mr/limiting_resource_adaptor.hpp>
 #include <shared_mutex>
 #include <unordered_map>
 
@@ -242,25 +240,10 @@ DI void sorted_insert(T* array, T item, int curr_size, int max_size)
 
 inline size_t get_device_memory_size()
 {
-  // Otherwise, we need to get the free memory from the device
   size_t free_mem, total_mem;
-  cudaMemGetInfo(&free_mem, &total_mem);
-
-  auto res = rmm::mr::get_current_device_resource();
-  auto limiting_adaptor =
-    dynamic_cast<rmm::mr::limiting_resource_adaptor<rmm::mr::cuda_async_memory_resource>*>(res);
-  // Did we specifiy an explicit memory limit?
-  if (limiting_adaptor) {
-    printf("limiting_adaptor->get_allocation_limit(): %fMiB\n",
-           limiting_adaptor->get_allocation_limit() / (double)1e6);
-    printf("used_mem: %fMiB\n", limiting_adaptor->get_allocated_bytes() / (double)1e6);
-    printf("free_mem: %fMiB\n",
-           (limiting_adaptor->get_allocation_limit() - limiting_adaptor->get_allocated_bytes()) /
-             (double)1e6);
-    return std::min(total_mem, limiting_adaptor->get_allocation_limit());
-  } else {
-    return total_mem;
-  }
+  RAFT_CUDA_TRY(cudaMemGetInfo(&free_mem, &total_mem));
+  // TODO (bdice): Restore limiting adaptor check after updating CCCL to support resource_cast
+  return total_mem;
 }
 
 }  // namespace cuopt

cpp/tests/mip/load_balancing_test.cu

@@ -32,7 +32,7 @@
 
 namespace cuopt::linear_programming::test {
 
-inline auto make_async() { return std::make_shared<rmm::mr::cuda_async_memory_resource>(); }
+inline auto make_async() { return rmm::mr::cuda_async_memory_resource(); }
 
 void init_handler(const raft::handle_t* handle_ptr)
 {
@@ -119,7 +119,7 @@ bounds_probe_results(detail::bound_presolve_t<int, double>& bnd_prb_0,
 void test_multi_probe(std::string path)
 {
   auto memory_resource = make_async();
-  rmm::mr::set_current_device_resource(memory_resource.get());
+  rmm::mr::set_current_device_resource(memory_resource);
   const raft::handle_t handle_{};
   cuopt::mps_parser::mps_data_model_t<int, double> mps_problem =
     cuopt::mps_parser::parse_mps<int, double>(path, false);