Bound CAGRA hash table sizing loops so they cannot spin forever

calc_hashmap_params() sizes its hash tables by growing a bit length
until the requested table fits, checking the supported maximum only
after the loop. hashmap::get_size() is `1U << bitlen`, which is
undefined once bitlen reaches 32 and in practice wraps back to a small
value, so a request needing more than 2^31 entries left the loop
condition permanently true. The bit length grew without bound, the
post-loop RAFT_EXPECTS was never reached, and the search hung on the
host instead of returning an error.

Bound each of the four sizing loops by the same maximum its post-loop
check already uses, so an oversized request exits the loop and raises
the existing error. The limits are hoisted into local constants so the
loop bound and the check cannot drift apart.

check_params() only caps itopk_size at 1024 for SINGLE_CTA, so
MULTI_CTA and MULTI_KERNEL were the reachable paths. Both are covered
by a new regression test; note that a regression there resurfaces as a
test timeout rather than a failed assertion.

Verified on an RTX PRO 6000: itopk_size 1,100,000,000 previously hung
indefinitely and now raises the hash_bitlen error in a few seconds,
while valid searches are unaffected. A standalone sweep of the old and
new loop logic over all four loops agrees on every in-range input.

Closes #2523

Signed-off-by: Shaunak Kapur <shaunakk@nvidia.com>

Author: shaunakkapur

cpp/src/neighbors/detail/cagra/search_plan.cuh

@@ -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
  */
 
@@ -265,12 +265,18 @@ struct search_plan_impl : public search_plan_impl_base {
       // table that each CTA has in the shared memory. This hash table is not
       // shared among CTAs. This hash table is reset and restored in each iteration.
       //
-      const uint32_t max_visited_nodes = mc_itopk_size + (graph_degree * 2);
-      small_hash_bitlen                = 8;  // 256
-      while (max_visited_nodes > hashmap::get_size(small_hash_bitlen) * max_fill_rate) {
+      const uint32_t max_visited_nodes       = mc_itopk_size + (graph_degree * 2);
+      constexpr size_t max_small_hash_bitlen = 14;  // 16K
+      small_hash_bitlen                      = 8;   // 256
+      // Stop at the supported maximum rather than growing without bound: hashmap::get_size()
+      // is 1U << bitlen, which is undefined (and wraps to a small value) once bitlen reaches
+      // 32, so an unbounded loop would never satisfy its exit condition again.
+      while (small_hash_bitlen <= max_small_hash_bitlen &&
+             max_visited_nodes > hashmap::get_size(small_hash_bitlen) * max_fill_rate) {
         small_hash_bitlen += 1;
       }
-      RAFT_EXPECTS(small_hash_bitlen <= 14, "small_hash_bitlen cannot be largen than 14 (16K)");
+      RAFT_EXPECTS(small_hash_bitlen <= max_small_hash_bitlen,
+                   "small_hash_bitlen cannot be largen than 14 (16K)");
       //
       // [traversed_hash_table]
       // Whether a node has ever been used as the starting point for a traversal
@@ -279,13 +285,15 @@ struct search_plan_impl : public search_plan_impl_base {
       //
       const auto max_traversed_nodes =
         mc_num_cta_per_query * max((size_t)mc_itopk_size, max_iterations);
-      unsigned min_bitlen = 11;  // 2K
+      unsigned min_bitlen                  = 11;  // 2K
+      constexpr int64_t max_hash_bitlen_mc = 25;  // 32M
       if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; }
       hash_bitlen = min_bitlen;
-      while (max_traversed_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
+      while (hash_bitlen <= max_hash_bitlen_mc &&
+             max_traversed_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
         hash_bitlen += 1;
       }
-      RAFT_EXPECTS(hash_bitlen <= 25, "hash_bitlen cannot be largen than 25 (32M)");
+      RAFT_EXPECTS(hash_bitlen <= max_hash_bitlen_mc, "hash_bitlen cannot be largen than 25 (32M)");
     } else {
       while (hashmap_mode == hash_mode::AUTO || hashmap_mode == hash_mode::SMALL) {
         //
@@ -300,7 +308,8 @@ struct search_plan_impl : public search_plan_impl_base {
         unsigned max_bitlen          = 13;  // 8K
         if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; }
         hash_bitlen = min_bitlen;
-        while (max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
+        while (hash_bitlen <= max_bitlen &&
+               max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
           hash_bitlen += 1;
         }
         if (hash_bitlen > max_bitlen) {
@@ -337,13 +346,15 @@ struct search_plan_impl : public search_plan_impl_base {
         // maximum fill rate of the hash table.
         //
         uint32_t max_visited_nodes = itopk_size + (search_width * graph_degree * max_iterations);
-        unsigned min_bitlen        = 11;  // 2K
+        unsigned min_bitlen        = 11;         // 2K
+        constexpr int64_t max_hash_bitlen = 20;  // 1M
         if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; }
         hash_bitlen = min_bitlen;
-        while (max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
+        while (hash_bitlen <= max_hash_bitlen &&
+               max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
           hash_bitlen += 1;
         }
-        RAFT_EXPECTS(hash_bitlen <= 20,
+        RAFT_EXPECTS(hash_bitlen <= max_hash_bitlen,
                      "hash_bitlen cannot be largen than 20 (1M). You can decrease itopk_size, "
                      "search_width or max_iterations to reduce the required hashmap size.");
       }

cpp/tests/CMakeLists.txt

@@ -211,6 +211,7 @@ ConfigureTest(
        neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu
        neighbors/ann_cagra/bug_iterative_cagra_build.cu
        neighbors/ann_cagra/bug_issue_93_reproducer.cu
+       neighbors/ann_cagra/bug_issue_2523_hashmap_bitlen.cu
   GPUS 1
   PERCENT 100
 )

cpp/tests/neighbors/ann_cagra/bug_issue_2523_hashmap_bitlen.cu

@@ -0,0 +1,135 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <gtest/gtest.h>
+
+#include "../cagra_padded_build_helpers.cuh"
+#include <cuvs/neighbors/cagra.hpp>
+
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_resources.hpp>
+#include <raft/core/error.hpp>
+#include <raft/random/rng.cuh>
+
+#include <cstdint>
+#include <utility>
+
+namespace cuvs::neighbors::cagra {
+
+/**
+ * @brief Regression test for issue #2523: hash table sizing must not loop forever.
+ *
+ * The hash tables in search_plan_impl::calc_hashmap_params() are sized by loops that grow
+ * a bit length until the requested table fits, with the supported maximum checked only
+ * after the loop. hashmap::get_size() is `1U << bitlen`, which is undefined once bitlen
+ * reaches 32 and in practice wraps back to a small value, so a request that needs a table
+ * larger than 2^31 entries made the loop condition permanently true. The bit length grew
+ * without bound, the post-loop RAFT_EXPECTS was never reached, and the search hung on the
+ * host instead of failing.
+ *
+ * check_params() only caps itopk_size at 1024 for SINGLE_CTA, so MULTI_CTA and MULTI_KERNEL
+ * are the two algorithms that can reach the sizing loops with an oversized request.
+ *
+ * These searches must now raise the existing "hash_bitlen cannot be larger than ..." error
+ * rather than hanging. A regression reappears as a test timeout, not a failed assertion.
+ */
+class cagra_hashmap_bitlen_no_hang_test : public ::testing::Test {
+ public:
+  using data_type  = float;
+  using index_type = uint32_t;
+
+ protected:
+  // Large enough that the traversed-node hash table would need more than 2^31 entries,
+  // which is what previously drove the sizing loop past the wrap point.
+  constexpr static size_t oversized_itopk = 1'100'000'000;
+  constexpr static size_t valid_itopk     = 64;
+
+  constexpr static int64_t n_dataset = 1000;
+  constexpr static int64_t n_dim     = 32;
+  constexpr static int64_t n_queries = 4;
+  constexpr static int64_t k         = 10;
+
+  void SetUp() override
+  {
+    dataset.emplace(raft::make_device_matrix<data_type, int64_t>(res, n_dataset, n_dim));
+    queries.emplace(raft::make_device_matrix<data_type, int64_t>(res, n_queries, n_dim));
+    neighbors.emplace(raft::make_device_matrix<index_type, int64_t>(res, n_queries, k));
+    distances.emplace(raft::make_device_matrix<data_type, int64_t>(res, n_queries, k));
+
+    raft::random::RngState r(1234ULL);
+    raft::random::uniform(
+      res, r, dataset->data_handle(), n_dataset * n_dim, data_type(-1), data_type(1));
+    raft::random::uniform(
+      res, r, queries->data_handle(), n_queries * n_dim, data_type(-1), data_type(1));
+
+    cagra::index_params index_params;
+    index_params.graph_degree              = 32;
+    index_params.intermediate_graph_degree = 64;
+
+    padded_.emplace(res, raft::make_const_mdspan(dataset->view()));
+    index_.emplace(cagra::build(res, index_params, padded_->view));
+    raft::resource::sync_stream(res);
+  }
+
+  void TearDown() override
+  {
+    index_.reset();
+    padded_.reset();
+    dataset.reset();
+    queries.reset();
+    neighbors.reset();
+    distances.reset();
+    raft::resource::sync_stream(res);
+  }
+
+  void search_with(cagra::search_algo algo, size_t itopk_size, uint32_t max_iterations)
+  {
+    cagra::search_params search_params;
+    search_params.algo           = algo;
+    search_params.itopk_size     = itopk_size;
+    search_params.search_width   = 8;
+    search_params.max_iterations = max_iterations;
+
+    cagra::search(res,
+                  search_params,
+                  *index_,
+                  raft::make_const_mdspan(queries->view()),
+                  neighbors->view(),
+                  distances->view());
+    raft::resource::sync_stream(res);
+  }
+
+  raft::resources res;
+  std::optional<cuvs::neighbors::test::padded_device_matrix_for_cagra<data_type>> padded_{};
+  std::optional<cagra::index<data_type, index_type>> index_         = std::nullopt;
+  std::optional<raft::device_matrix<data_type, int64_t>> dataset    = std::nullopt;
+  std::optional<raft::device_matrix<data_type, int64_t>> queries    = std::nullopt;
+  std::optional<raft::device_matrix<index_type, int64_t>> neighbors = std::nullopt;
+  std::optional<raft::device_matrix<data_type, int64_t>> distances  = std::nullopt;
+};
+
+// MULTI_CTA sizes the shared traversed-node table from itopk_size and search_width.
+// This is the path confirmed to hang before the fix.
+TEST_F(cagra_hashmap_bitlen_no_hang_test, MultiCtaOversizedItopkThrows)
+{
+  EXPECT_THROW(search_with(cagra::search_algo::MULTI_CTA, oversized_itopk, 0), raft::exception);
+}
+
+// MULTI_KERNEL reaches the small-hash loop first and then the normal-hash loop.
+// max_iterations is pinned so the sizing inputs do not depend on the auto-derived value.
+TEST_F(cagra_hashmap_bitlen_no_hang_test, MultiKernelOversizedItopkThrows)
+{
+  EXPECT_THROW(search_with(cagra::search_algo::MULTI_KERNEL, oversized_itopk, 32), raft::exception);
+}
+
+// Guards against the bound rejecting requests that were previously accepted.
+TEST_F(cagra_hashmap_bitlen_no_hang_test, ValidItopkStillSucceeds)
+{
+  EXPECT_NO_THROW(search_with(cagra::search_algo::MULTI_CTA, valid_itopk, 0));
+  EXPECT_NO_THROW(search_with(cagra::search_algo::MULTI_KERNEL, valid_itopk, 0));
+  EXPECT_NO_THROW(search_with(cagra::search_algo::SINGLE_CTA, valid_itopk, 0));
+}
+
+}  // namespace cuvs::neighbors::cagra