From 319920ec16aeb7cf3a5851034c45213d362edf68 Mon Sep 17 00:00:00 2001
From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com>
Date: Fri, 15 May 2026 22:45:51 +0200
Subject: [PATCH] Persistent CAGRA: benchmark group and bad config warnings
 (#2091)

Improve user experience with CAGRA search persistent kernel mode:

- make sure SINGLE_CTA algo is selected in persistent mode if it is not passed explicitly
- warn about incompatible configuration
- add a cuvs benchmarking group for reference

Authors:
  - Artem M. Chirkin (https://github.com/achirkin)

Approvers:
  - Tamas Bela Feher (https://github.com/tfeher)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: https://github.com/rapidsai/cuvs/pull/2091
---
 cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h   | 29 +++++++++++++++++++
 .../neighbors/detail/cagra/search_plan.cuh    | 13 ++++++++-
 .../cuvs_bench/config/algos/cuvs_cagra.yaml   | 12 ++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
index 98dd94c2e1..87111e4761 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
@@ -27,6 +27,7 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <atomic>
 #include <cassert>
 #include <fstream>
 #include <iostream>
@@ -35,11 +36,37 @@
 #include <raft/util/integer_utils.hpp>
 #include <stdexcept>
 #include <string>
+#include <thread>
 #include <type_traits>
 #include <vector>
 
 namespace cuvs::bench {
 
+namespace detail {
+
+/** If persistent CAGRA search uses few benchmark threads, log a throughput hint once per process.
+ */
+inline void maybe_log_cagra_persistent_concurrency_hint(bool persistent_search)
+{
+  if (!persistent_search) { return; }
+
+  const unsigned hc = std::max(1u, static_cast<unsigned>(std::thread::hardware_concurrency()));
+  const unsigned bn = static_cast<unsigned>(std::max(0, benchmark_n_threads));
+  if (bn >= 2u * hc) { return; }
+
+  static std::atomic<bool> logged{false};
+  bool expected = false;
+  if (!logged.compare_exchange_strong(expected, true)) { return; }
+
+  const unsigned threads_rec = 16u * hc;
+  RAFT_LOG_INFO(
+    "CAGRA persistent search benefits from high client concurrency (try `--mode=throughput "
+    "--threads=1:%u`).",
+    threads_rec);
+}
+
+}  // namespace detail
+
 enum class AllocatorType { kHostPinned, kHostHugePage, kDevice };
 enum class CagraBuildAlgo { kAuto, kIvfPq, kNnDescent };
 enum class CagraMergeType { kPhysical, kLogical };
@@ -325,6 +352,8 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
   } else {
     if (dynamic_batcher_) { dynamic_batcher_.reset(); }
   }
+
+  detail::maybe_log_cagra_persistent_concurrency_hint(search_params_.persistent);
 }
 
 template <typename T, typename IdxT>
diff --git a/cpp/src/neighbors/detail/cagra/search_plan.cuh b/cpp/src/neighbors/detail/cagra/search_plan.cuh
index 68f22d7688..8a9d79e177 100644
--- a/cpp/src/neighbors/detail/cagra/search_plan.cuh
+++ b/cpp/src/neighbors/detail/cagra/search_plan.cuh
@@ -13,6 +13,7 @@
 // #include "topk_for_cagra/topk.h"
 
 #include <raft/core/device_mdspan.hpp>
+#include <raft/core/logger.hpp>
 #include <raft/core/resources.hpp>
 
 #include <cuvs/distance/distance.hpp>
@@ -108,7 +109,17 @@ struct search_plan_impl_base : public search_params {
       graph_degree(graph_degree),
       topk(topk)
   {
-    if (algo == search_algo::AUTO) {
+    if (persistent) {
+      if (algo == search_algo::AUTO) {
+        algo = search_algo::SINGLE_CTA;
+        RAFT_LOG_DEBUG("Auto strategy: persistent enabled, selecting single-cta");
+      } else if (algo != search_algo::SINGLE_CTA) {
+        RAFT_LOG_WARN(
+          "CAGRA persistent kernel is only supported with SINGLE_CTA search (algo=%d); "
+          "persistent will have no effect",
+          static_cast<int>(algo));
+      }
+    } else if (algo == search_algo::AUTO) {
       const size_t num_sm = raft::getMultiProcessorCount();
       if (itopk_size <= 512 && search_params::max_queries >= num_sm * 2lu) {
         algo = search_algo::SINGLE_CTA;
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml
index 240d2ac8ee..cf59c3ce22 100644
--- a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml
+++ b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml
@@ -19,3 +19,15 @@ groups:
     search:
       itopk: [32]
       search_width: [1, 2]
+  persistent:
+    build:
+      graph_degree: [32, 64, 96]
+      intermediate_graph_degree: [128]
+      graph_build_algo: ["IVF_PQ"]
+    search:
+      persistent: [true]
+      persistent_device_usage: [0.95]
+      algo: ["single_cta"]
+      itopk: [32, 64, 128, 256, 512]
+      max_iterations: [0, 16]
+      search_width: [1, 2, 4, 8]