From 2fb6db47804d2c38bf51260f8fcc329ae13027d5 Mon Sep 17 00:00:00 2001
From: tarangj <jaintarang2015@gmail.com>
Date: Fri, 28 Aug 2026 16:36:53 -0700
Subject: [PATCH 1/8] Use async resources for prefetch benchmarks

---
 cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h    | 12 ++++++++++--
 cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h      |  2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h   |  2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h     |  2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h |  2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h     |  2 +-
 6 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
index 1a276e8cc8..74bc60d8f8 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
@@ -18,10 +18,11 @@
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
+#include <rmm/cuda_stream_pool.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/mr/cuda_async_managed_memory_resource.hpp>
 #include <rmm/mr/failure_callback_resource_adaptor.hpp>
-#include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>
@@ -31,6 +32,13 @@
 
 namespace cuvs::bench {
 
+/** Create streams that do not implicitly synchronize with the default stream. */
+inline auto make_non_blocking_stream_pool(size_t n_streams)
+  -> std::shared_ptr<rmm::cuda_stream_pool>
+{
+  return std::make_shared<rmm::cuda_stream_pool>(n_streams, rmm::cuda_stream::flags::non_blocking);
+}
+
 inline auto parse_metric_type(cuvs::bench::Metric metric) -> cuvs::distance::DistanceType
 {
   if (metric == cuvs::bench::Metric::kInnerProduct) {
@@ -67,7 +75,7 @@ inline auto rmm_oom_callback(std::size_t bytes, void*) -> bool
  */
 class shared_raft_resources {
  public:
-  using large_mr_type = rmm::mr::managed_memory_resource;
+  using large_mr_type = rmm::mr::cuda_async_managed_memory_resource;
 
   shared_raft_resources()
   try : large_mr_() {
diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
index ed067fac39..a076021c46 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
@@ -679,7 +679,7 @@ void cuvs_cagra<T, IdxT>::search_base(
         raft::resources composite_handle(handle_);
         size_t n_streams = cagra_indices.size();
         raft::resource::set_cuda_stream_pool(composite_handle,
-                                             std::make_shared<rmm::cuda_stream_pool>(n_streams));
+                                             make_non_blocking_stream_pool(n_streams));
 
         cuvs::neighbors::composite::composite_index<T, IdxT, algo_base::index_type> composite(
           cagra_indices);
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
index ae334f70d7..1d92c688c8 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
@@ -90,7 +90,7 @@ void cuvs_ivf_flat<T, IdxT>::build(const T* dataset, size_t nrow)
 {
   // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping.
   size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, std::make_shared<rmm::cuda_stream_pool>(n_streams));
+  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   index_ = std::make_shared<cuvs::neighbors::ivf_flat::index<T, IdxT>>(
     std::move(cuvs::neighbors::ivf_flat::build(
       handle_,
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
index 161563d8fe..b8b1c49f79 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
@@ -120,7 +120,7 @@ void cuvs_ivf_pq<T, IdxT>::build(const T* dataset, size_t nrow)
 {
   // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping.
   size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, std::make_shared<rmm::cuda_stream_pool>(n_streams));
+  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   auto dataset_v = raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dim_);
   std::make_shared<cuvs::neighbors::ivf_pq::index<IdxT>>(
     std::move(cuvs::neighbors::ivf_pq::build(handle_, index_params_, dataset_v)))
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
index ca8f77b808..5198d1c84b 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
@@ -100,7 +100,7 @@ void cuvs_ivf_rabitq<T, IdxT>::build(const T* dataset, size_t nrow)
 {
   // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping.
   size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, std::make_shared<rmm::cuda_stream_pool>(n_streams));
+  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   auto dataset_v = raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dim_);
   std::make_shared<cuvs::neighbors::ivf_rabitq::index<IdxT>>(
     std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v)))
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
index 1503e6bb84..42adad542b 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
@@ -84,7 +84,7 @@ template <typename T>
 void cuvs_ivf_sq<T>::build(const T* dataset, size_t nrow)
 {
   size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, std::make_shared<rmm::cuda_stream_pool>(n_streams));
+  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   index_ = std::make_shared<cuvs::neighbors::ivf_sq::index<uint8_t>>(
     std::move(cuvs::neighbors::ivf_sq::build(
       handle_,

From 4a7f4f5e6c34e04a3dd9f8b18cd533dd855d3b8c Mon Sep 17 00:00:00 2001
From: tarangj <jaintarang2015@gmail.com>
Date: Fri, 28 Aug 2026 17:06:18 -0700
Subject: [PATCH 2/8] style

---
 cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h    | 2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h   | 2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h     | 2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h     | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
index 74bc60d8f8..bb0376a6af 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
index 1d92c688c8..68f8f95481 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
index b8b1c49f79..b4a81714c6 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
index 5198d1c84b..7efb43e3dc 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
index 42adad542b..787e945f71 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once

From fd4f3f11773a852332787fec5772e78bd37709b7 Mon Sep 17 00:00:00 2001
From: tarangj <jaintarang2015@gmail.com>
Date: Tue, 1 Sep 2026 09:18:52 -0700
Subject: [PATCH 3/8] address pr reviews

---
 cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h    | 10 +++-------
 cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h      |  3 ++-
 cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h   |  3 ---
 cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h     |  3 ---
 cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h |  3 ---
 cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h     |  2 --
 6 files changed, 5 insertions(+), 19 deletions(-)

diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
index bb0376a6af..8c785cf368 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
@@ -15,6 +15,7 @@
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/cuda_stream_pool.hpp>
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
@@ -32,13 +33,6 @@
 
 namespace cuvs::bench {
 
-/** Create streams that do not implicitly synchronize with the default stream. */
-inline auto make_non_blocking_stream_pool(size_t n_streams)
-  -> std::shared_ptr<rmm::cuda_stream_pool>
-{
-  return std::make_shared<rmm::cuda_stream_pool>(n_streams, rmm::cuda_stream::flags::non_blocking);
-}
-
 inline auto parse_metric_type(cuvs::bench::Metric metric) -> cuvs::distance::DistanceType
 {
   if (metric == cuvs::bench::Metric::kInnerProduct) {
@@ -140,6 +134,8 @@ class configured_raft_resources {
   /** Default constructor creates all resources anew. */
   configured_raft_resources() : configured_raft_resources{std::make_shared<shared_raft_resources>()}
   {
+    raft::resource::set_cuda_stream_pool(
+      *res_, std::make_shared<rmm::cuda_stream_pool>(1, rmm::cuda_stream::flags::non_blocking));
   }
 
   configured_raft_resources(configured_raft_resources&&);
diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
index a076021c46..7e95d181c4 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
@@ -679,7 +679,8 @@ void cuvs_cagra<T, IdxT>::search_base(
         raft::resources composite_handle(handle_);
         size_t n_streams = cagra_indices.size();
         raft::resource::set_cuda_stream_pool(composite_handle,
-                                             make_non_blocking_stream_pool(n_streams));
+                                             std::make_shared<rmm::cuda_stream_pool>(
+                                               n_streams, rmm::cuda_stream::flags::non_blocking));
 
         cuvs::neighbors::composite::composite_index<T, IdxT, algo_base::index_type> composite(
           cagra_indices);
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
index 68f8f95481..9d73990986 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
@@ -88,9 +88,6 @@ class cuvs_ivf_flat : public algo<T>, public algo_gpu {
 template <typename T, typename IdxT>
 void cuvs_ivf_flat<T, IdxT>::build(const T* dataset, size_t nrow)
 {
-  // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping.
-  size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   index_ = std::make_shared<cuvs::neighbors::ivf_flat::index<T, IdxT>>(
     std::move(cuvs::neighbors::ivf_flat::build(
       handle_,
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
index b4a81714c6..0ddb41d19e 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
@@ -118,9 +118,6 @@ void cuvs_ivf_pq<T, IdxT>::load(const std::string& file)
 template <typename T, typename IdxT>
 void cuvs_ivf_pq<T, IdxT>::build(const T* dataset, size_t nrow)
 {
-  // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping.
-  size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   auto dataset_v = raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dim_);
   std::make_shared<cuvs::neighbors::ivf_pq::index<IdxT>>(
     std::move(cuvs::neighbors::ivf_pq::build(handle_, index_params_, dataset_v)))
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
index 7efb43e3dc..d9b12da607 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
@@ -98,9 +98,6 @@ void cuvs_ivf_rabitq<T, IdxT>::load(const std::string& file)
 template <typename T, typename IdxT>
 void cuvs_ivf_rabitq<T, IdxT>::build(const T* dataset, size_t nrow)
 {
-  // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping.
-  size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   auto dataset_v = raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dim_);
   std::make_shared<cuvs::neighbors::ivf_rabitq::index<IdxT>>(
     std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v)))
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
index 787e945f71..8a96951803 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
@@ -83,8 +83,6 @@ class cuvs_ivf_sq : public algo<T>, public algo_gpu {
 template <typename T>
 void cuvs_ivf_sq<T>::build(const T* dataset, size_t nrow)
 {
-  size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, make_non_blocking_stream_pool(n_streams));
   index_ = std::make_shared<cuvs::neighbors::ivf_sq::index<uint8_t>>(
     std::move(cuvs::neighbors::ivf_sq::build(
       handle_,

From 39100a2587b7033651f393b180df8acaf4ea56f0 Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Wed, 2 Sep 2026 11:40:16 -0700
Subject: [PATCH 4/8] pdate behavior

---
 cpp/bench/ann/src/common/util.hpp             | 24 -------------------
 cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h |  3 +--
 cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h   |  3 +--
 3 files changed, 2 insertions(+), 28 deletions(-)

diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp
index 6e4c57b35e..cb23206e11 100644
--- a/cpp/bench/ann/src/common/util.hpp
+++ b/cpp/bench/ann/src/common/util.hpp
@@ -144,30 +144,6 @@ struct cuda_timer {
   }
 };
 
-#ifndef BUILD_CPU_ONLY
-// ATM, rmm::stream does not support passing in flags; hence this helper type.
-struct non_blocking_stream {
-  non_blocking_stream() { cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking); }
-  ~non_blocking_stream() noexcept
-  {
-    if (stream_ != nullptr) { cudaStreamDestroy(stream_); }
-  }
-  non_blocking_stream(non_blocking_stream const&) = delete;
-  non_blocking_stream(non_blocking_stream&& other) noexcept { std::swap(stream_, other.stream_); }
-  auto operator=(non_blocking_stream const&) -> non_blocking_stream& = delete;
-  auto operator=(non_blocking_stream&&) -> non_blocking_stream&      = delete;
-  [[nodiscard]] auto view() const noexcept -> cudaStream_t { return stream_; }
-
- private:
-  cudaStream_t stream_{nullptr};
-};
-
-namespace detail {
-inline std::vector<non_blocking_stream> global_stream_pool(0);
-inline std::mutex gsp_mutex;
-}  // namespace detail
-#endif
-
 /**
  * Get a stream associated with the current benchmark thread.
  *
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
index 8c785cf368..0ca96cb3c0 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
@@ -134,8 +134,7 @@ class configured_raft_resources {
   /** Default constructor creates all resources anew. */
   configured_raft_resources() : configured_raft_resources{std::make_shared<shared_raft_resources>()}
   {
-    raft::resource::set_cuda_stream_pool(
-      *res_, std::make_shared<rmm::cuda_stream_pool>(1, rmm::cuda_stream::flags::non_blocking));
+    raft::resource::set_cuda_stream_pool(*res_, std::make_shared<rmm::cuda_stream_pool>(1));
   }
 
   configured_raft_resources(configured_raft_resources&&);
diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
index 7e95d181c4..ed067fac39 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
@@ -679,8 +679,7 @@ void cuvs_cagra<T, IdxT>::search_base(
         raft::resources composite_handle(handle_);
         size_t n_streams = cagra_indices.size();
         raft::resource::set_cuda_stream_pool(composite_handle,
-                                             std::make_shared<rmm::cuda_stream_pool>(
-                                               n_streams, rmm::cuda_stream::flags::non_blocking));
+                                             std::make_shared<rmm::cuda_stream_pool>(n_streams));
 
         cuvs::neighbors::composite::composite_index<T, IdxT, algo_base::index_type> composite(
           cagra_indices);

From 5e203e0e5dca65f87f2302df2cbae98e91d6bee0 Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Wed, 2 Sep 2026 11:50:10 -0700
Subject: [PATCH 5/8] style

---
 cpp/bench/ann/src/common/util.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp
index cb23206e11..bb9a1dead7 100644
--- a/cpp/bench/ann/src/common/util.hpp
+++ b/cpp/bench/ann/src/common/util.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once

From 6a637e5108791c79bc2a6c4d52283536fd02b332 Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Wed, 2 Sep 2026 22:23:49 -0700
Subject: [PATCH 6/8] fix compilation

---
 cpp/bench/ann/src/common/util.hpp | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp
index bb9a1dead7..069d257717 100644
--- a/cpp/bench/ann/src/common/util.hpp
+++ b/cpp/bench/ann/src/common/util.hpp
@@ -7,6 +7,10 @@
 #include "ann_types.hpp"
 #include "cuda_stub.hpp"  // cuda-related utils
 
+#ifndef BUILD_CPU_ONLY
+#include <rmm/cuda_stream.hpp>
+#endif
+
 #if __has_include(<nvtx3/nvToolsExt.h>)
 #define ANN_BENCH_NVTX3_HEADERS_FOUND
 #include <nvtx3/nvToolsExt.h>
@@ -144,6 +148,13 @@ struct cuda_timer {
   }
 };
 
+#ifndef BUILD_CPU_ONLY
+namespace detail {
+inline std::vector<rmm::cuda_stream> global_stream_pool(0);
+inline std::mutex gsp_mutex;
+}  // namespace detail
+#endif
+
 /**
  * Get a stream associated with the current benchmark thread.
  *
@@ -156,7 +167,9 @@ inline auto get_stream_from_global_pool() -> cudaStream_t
 #ifndef BUILD_CPU_ONLY
   std::lock_guard guard(detail::gsp_mutex);
   if (static_cast<int>(detail::global_stream_pool.size()) < benchmark_n_threads) {
-    detail::global_stream_pool.resize(benchmark_n_threads);
+    while (static_cast<int>(detail::global_stream_pool.size()) < benchmark_n_threads) {
+      detail::global_stream_pool.emplace_back(rmm::cuda_stream::flags::non_blocking);
+    }
   }
   return detail::global_stream_pool[benchmark_thread_id].view();
 #else

From 0e26ee5d2835aa83998d31a85ce6a20cdbe0e6db Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Thu, 3 Sep 2026 13:49:17 -0700
Subject: [PATCH 7/8] try linking rmm

---
 cpp/bench/ann/CMakeLists.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt
index 90d23d9aef..786d45ccc1 100644
--- a/cpp/bench/ann/CMakeLists.txt
+++ b/cpp/bench/ann/CMakeLists.txt
@@ -181,6 +181,7 @@ function(ConfigureAnnBench)
             sqlite3
             Threads::Threads
             $<$<BOOL:${GPU_BUILD}>:CUDA::cudart_static>
+            $<$<BOOL:${GPU_BUILD}>:rmm::rmm>
             $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
             $<TARGET_NAME_IF_EXISTS:conda_env>
   )

From b7a0a65dfad98f0408b346b7a6966fd1ee3c18e1 Mon Sep 17 00:00:00 2001
From: tarang-jain <tarangj@nvidia.com>
Date: Thu, 3 Sep 2026 14:08:36 -0700
Subject: [PATCH 8/8] move to upper constructor

---
 cpp/bench/ann/CMakeLists.txt                     | 2 +-
 cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h    | 2 +-
 cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h   | 1 -
 cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h     | 1 -
 cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 1 -
 cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h     | 2 --
 6 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt
index 786d45ccc1..50bfa3a1e7 100644
--- a/cpp/bench/ann/CMakeLists.txt
+++ b/cpp/bench/ann/CMakeLists.txt
@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
index 0ca96cb3c0..539cc37a30 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
@@ -129,12 +129,12 @@ class configured_raft_resources {
   {
     raft::resource::set_large_workspace_resource(
       *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()});
+    raft::resource::set_cuda_stream_pool(*res_, std::make_shared<rmm::cuda_stream_pool>(1));
   }
 
   /** Default constructor creates all resources anew. */
   configured_raft_resources() : configured_raft_resources{std::make_shared<shared_raft_resources>()}
   {
-    raft::resource::set_cuda_stream_pool(*res_, std::make_shared<rmm::cuda_stream_pool>(1));
   }
 
   configured_raft_resources(configured_raft_resources&&);
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
index 9d73990986..8ed5adee26 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h
@@ -15,7 +15,6 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/unary_op.cuh>
 #include <raft/util/cudart_utils.hpp>
-#include <rmm/cuda_stream_pool.hpp>
 
 #include <cassert>
 #include <fstream>
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
index 0ddb41d19e..363d065e91 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h
@@ -20,7 +20,6 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/unary_op.cuh>
 #include <raft/util/cudart_utils.hpp>
-#include <rmm/cuda_stream_pool.hpp>
 
 #include <type_traits>
 
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
index d9b12da607..0c41ee1c96 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h
@@ -18,7 +18,6 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/unary_op.cuh>
 #include <raft/util/cudart_utils.hpp>
-#include <rmm/cuda_stream_pool.hpp>
 
 #include <type_traits>
 
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
index 8a96951803..5bf0098eaa 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h
@@ -12,9 +12,7 @@
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/device_resources.hpp>
 #include <raft/core/logger.hpp>
-#include <raft/core/resource/cuda_stream.hpp>
 #include <raft/util/cudart_utils.hpp>
-#include <rmm/cuda_stream_pool.hpp>
 
 #include <cassert>
 #include <memory>