Fast CAGRA Index Merge

c/include/cuvs/neighbors/cagra.h

@@ -228,6 +228,27 @@ struct cuvsCagraIndexParams {
 
 typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t;
 
+/** Algorithm used to merge physical CAGRA indices. */
+enum cuvsCagraMergeAlgo {
+  CUVS_CAGRA_MERGE_AUTO     = 0,
+  CUVS_CAGRA_MERGE_FASTENER = 1,
+  CUVS_CAGRA_MERGE_REBUILD  = 2
+};
+
+/** Parameters controlling how physical CAGRA indices are merged. */
+struct cuvsCagraMergeParams {
+  enum cuvsCagraMergeAlgo algo;
+  uint32_t levels;
+  uint32_t root_fanout;
+  uint32_t lower_fanout;
+  double leader_fraction;
+  uint32_t max_leaders;
+  uint32_t leaf_size;
+  uint32_t leaf_degree;
+};
+
+typedef struct cuvsCagraMergeParams* cuvsCagraMergeParams_t;
+
 /**
  * @brief Allocate CAGRA Index params, and populate with default values
  *
@@ -244,6 +265,12 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* param
  */
 CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params);
 
+/** Allocate CAGRA merge params and populate them with AUTO defaults. */
+CUVS_EXPORT cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params);
+
+/** De-allocate CAGRA merge params. */
+CUVS_EXPORT cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params);
+
 /**
  * @brief Allocate CAGRA Compression params, and populate with default values
  *
@@ -907,7 +934,7 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res,
  *
  * All input indices must have been built with the same data type (`index.dtype`) and
  * have the same dimensionality (`index.dims`). The merged index uses the output
- * parameters specified in `cuvsCagraIndexParams`.
+ * parameters specified in `cuvsCagraIndexParams`. The merge algorithm is selected automatically.
  *
  * Input indices must have:
  *  - `index.dtype.code` and `index.dtype.bits` matching across all indices.
@@ -947,7 +974,7 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res,
  * @endcode
  *
  * @param[in] res cuvsResources_t opaque C handle
- * @param[in] params cuvsCagraIndexParams_t parameters controlling merge behavior
+ * @param[in] params cuvsCagraIndexParams_t parameters for the output index
  * @param[in] indices Array of input cuvsCagraIndex_t handles to merge
  * @param[in] num_indices Number of input indices
  * @param[in] filter Filter that can be used to filter out vectors from the merged index
@@ -961,6 +988,26 @@ CUVS_EXPORT cuvsError_t cuvsCagraMerge(cuvsResources_t res,
                            cuvsFilter filter,
                            cuvsCagraIndex_t output_index);
 
+/**
+ * @brief Merge multiple CAGRA indices with explicit merge parameters.
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] params cuvsCagraIndexParams_t parameters for the output index
+ * @param[in] merge_params cuvsCagraMergeParams_t parameters controlling the merge algorithm, or
+ *                         NULL to use AUTO defaults
+ * @param[in] indices Array of input cuvsCagraIndex_t handles to merge
+ * @param[in] num_indices Number of input indices
+ * @param[in] filter Filter that can be used to filter out vectors from the merged index
+ * @param[out] output_index Output handle initialized with `cuvsCagraIndexCreate`
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res,
+                                                 cuvsCagraIndexParams_t params,
+                                                 cuvsCagraMergeParams_t merge_params,
+                                                 cuvsCagraIndex_t* indices,
+                                                 size_t num_indices,
+                                                 cuvsFilter filter,
+                                                 cuvsCagraIndex_t output_index);
+
 /**
  * @}
  */

c/src/neighbors/cagra.cpp

@@ -384,6 +384,7 @@ void* _deserialize(cuvsResources_t res, const char* filename)
 template <typename T>
 void* _merge(cuvsResources_t res,
              cuvsCagraIndexParams params,
+             const cuvs::neighbors::cagra::merge_params& merge_params,
              cuvsCagraIndex_t* indices,
              size_t num_indices,
              cuvsFilter filter)
@@ -425,16 +426,16 @@ void* _merge(cuvsResources_t res,
 
   if (filter.type == NO_FILTER) {
     return new cuvs::neighbors::cagra::index<T, uint32_t>(
-      cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs));
+      cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs, merge_params));
   } else if (filter.type == BITSET) {
     using filter_mdspan_type    = raft::device_vector_view<std::uint32_t, int64_t, raft::row_major>;
     auto removed_indices_tensor = reinterpret_cast<DLManagedTensor*>(filter.addr);
     auto removed_indices = cuvs::core::from_dlpack<filter_mdspan_type>(removed_indices_tensor);
     cuvs::core::bitset_view<std::uint32_t, int64_t> removed_indices_bitset(
       removed_indices, total_size);
     auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset);
-    return new cuvs::neighbors::cagra::index<T, uint32_t>(
-      cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs, bitset_filter_obj));
+    return new cuvs::neighbors::cagra::index<T, uint32_t>(cuvs::neighbors::cagra::merge(
+      *res_ptr, params_cpp, index_ptrs, merge_params, bitset_filter_obj));
   } else {
     RAFT_FAIL("Unsupported filter type: BITMAP");
   }
@@ -834,34 +835,58 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res,
                                       size_t num_indices,
                                       cuvsFilter filter,
                                       cuvsCagraIndex_t output_index)
+{
+  return cuvsCagraMergeWithParams(
+    res, params, nullptr, indices, num_indices, filter, output_index);
+}
+
+extern "C" cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res,
+                                                cuvsCagraIndexParams_t params,
+                                                cuvsCagraMergeParams_t merge_params,
+                                                cuvsCagraIndex_t* indices,
+                                                size_t num_indices,
+                                                cuvsFilter filter,
+                                                cuvsCagraIndex_t output_index)
 {
   return cuvs::core::translate_exceptions([=] {
-    // Basic checks on inputs
     RAFT_EXPECTS(indices != nullptr && num_indices > 0, "indices array cannot be null or empty");
     RAFT_EXPECTS(params != nullptr, "params cannot be null");
 
-    // Use first index dtype as reference
+    auto merge_params_cpp = cuvs::neighbors::cagra::merge_params{};
+    if (merge_params != nullptr) {
+      RAFT_EXPECTS(merge_params->algo >= CUVS_CAGRA_MERGE_AUTO &&
+                     merge_params->algo <= CUVS_CAGRA_MERGE_REBUILD,
+                   "Unsupported CAGRA merge algorithm");
+      merge_params_cpp = {
+        .algo = static_cast<cuvs::neighbors::cagra::merge_algo>(merge_params->algo),
+        .levels = merge_params->levels,
+        .root_fanout = merge_params->root_fanout,
+        .lower_fanout = merge_params->lower_fanout,
+        .leader_fraction = merge_params->leader_fraction,
+        .max_leaders = merge_params->max_leaders,
+        .leaf_size = merge_params->leaf_size,
+        .leaf_degree = merge_params->leaf_degree};
+    }
     auto dtype = (*indices[0]).dtype;
     for (size_t i = 1; i < num_indices; ++i) {
       RAFT_EXPECTS((*indices[i]).dtype.code == dtype.code && (*indices[i]).dtype.bits == dtype.bits,
                    "All input indices must have the same data type");
       RAFT_EXPECTS((*indices[i]).addr != 0, "All input indices must be built (non-empty)");
     }
     RAFT_EXPECTS(output_index != nullptr, "Output index pointer must not be null");
-    output_index->dtype = dtype;  // output index type matches inputs
-    // Dispatch based on data type
+    output_index->dtype = dtype;
     if (dtype.code == kDLFloat && dtype.bits == 32) {
-      output_index->addr =
-        reinterpret_cast<uintptr_t>(_merge<float>(res, *params, indices, num_indices, filter));
+      output_index->addr = reinterpret_cast<uintptr_t>(
+        _merge<float>(res, *params, merge_params_cpp, indices, num_indices, filter));
     } else if (dtype.code == kDLFloat && dtype.bits == 16) {
-      output_index->addr =
-        reinterpret_cast<uintptr_t>(_merge<half>(res, *params, indices, num_indices, filter));
+      output_index->addr = reinterpret_cast<uintptr_t>(
+        _merge<half>(res, *params, merge_params_cpp, indices, num_indices, filter));
     } else if (dtype.code == kDLInt && dtype.bits == 8) {
-      output_index->addr =
-        reinterpret_cast<uintptr_t>(_merge<int8_t>(res, *params, indices, num_indices, filter));
+      output_index->addr = reinterpret_cast<uintptr_t>(
+        _merge<int8_t>(res, *params, merge_params_cpp, indices, num_indices, filter));
     } else if (dtype.code == kDLUInt && dtype.bits == 8) {
-      output_index->addr =
-        reinterpret_cast<uintptr_t>(_merge<uint8_t>(res, *params, indices, num_indices, filter));
+      output_index->addr = reinterpret_cast<uintptr_t>(
+        _merge<uint8_t>(res, *params, merge_params_cpp, indices, num_indices, filter));
     } else {
       RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits);
     }
@@ -907,6 +932,26 @@ extern "C" cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params
   });
 }
 
+extern "C" cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto defaults = cuvs::neighbors::cagra::merge_params{};
+    *params       = new cuvsCagraMergeParams{.algo            = CUVS_CAGRA_MERGE_AUTO,
+                                             .levels          = defaults.levels,
+                                             .root_fanout     = defaults.root_fanout,
+                                             .lower_fanout    = defaults.lower_fanout,
+                                             .leader_fraction = defaults.leader_fraction,
+                                             .max_leaders     = defaults.max_leaders,
+                                             .leaf_size       = defaults.leaf_size,
+                                             .leaf_degree     = defaults.leaf_degree};
+  });
+}
+
+extern "C" cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params)
+{
+  return cuvs::core::translate_exceptions([=] { delete params; });
+}
+
 extern "C" cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* params)
 {
   return cuvs::core::translate_exceptions([=] {

c/tests/neighbors/ann_cagra_c.cu

@@ -511,7 +511,13 @@ TEST(CagraC, BuildMergeSearch)
   filter.addr = 0;
 
   cuvsCagraIndex_t index_array[2] = {index_main, index_add};
-  ASSERT_EQ(cuvsCagraMerge(res, build_params, index_array, 2, filter, index_merged), CUVS_SUCCESS);
+  cuvsCagraMergeParams_t merge_params;
+  ASSERT_EQ(cuvsCagraMergeParamsCreate(&merge_params), CUVS_SUCCESS);
+  EXPECT_EQ(merge_params->algo, CUVS_CAGRA_MERGE_AUTO);
+  merge_params->algo = CUVS_CAGRA_MERGE_REBUILD;
+  ASSERT_EQ(
+    cuvsCagraMergeWithParams(res, build_params, merge_params, index_array, 2, filter, index_merged),
+    CUVS_SUCCESS);
 
   int64_t merged_dim = -1;
   ASSERT_EQ(cuvsCagraIndexGetDims(index_merged, &merged_dim), CUVS_SUCCESS);
@@ -562,6 +568,7 @@ TEST(CagraC, BuildMergeSearch)
   EXPECT_NEAR(distance_host, 0.0f, 1e-6);
 
   cuvsCagraSearchParamsDestroy(search_params);
+  cuvsCagraMergeParamsDestroy(merge_params);
   cuvsCagraIndexParamsDestroy(build_params);
   cuvsCagraIndexDestroy(index_merged);
   cuvsCagraIndexDestroy(index_add);

cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.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
  */
 
@@ -423,6 +423,37 @@ void parse_build_param(const nlohmann::json& conf,
       throw std::runtime_error("invalid value for merge_type");
     }
   }
+  if (conf.contains("merge_algo")) {
+    std::string algo = conf.at("merge_algo");
+    if (algo == "AUTO") {
+      param.merge_params.algo = cuvs::neighbors::cagra::merge_algo::AUTO;
+    } else if (algo == "FASTENER") {
+      param.merge_params.algo = cuvs::neighbors::cagra::merge_algo::FASTENER;
+    } else if (algo == "REBUILD") {
+      param.merge_params.algo = cuvs::neighbors::cagra::merge_algo::REBUILD;
+    } else {
+      throw std::runtime_error("invalid value for merge_algo");
+    }
+  }
+  if (conf.contains("fastener_levels")) { param.merge_params.levels = conf.at("fastener_levels"); }
+  if (conf.contains("fastener_root_fanout")) {
+    param.merge_params.root_fanout = conf.at("fastener_root_fanout");
+  }
+  if (conf.contains("fastener_lower_fanout")) {
+    param.merge_params.lower_fanout = conf.at("fastener_lower_fanout");
+  }
+  if (conf.contains("fastener_leader_fraction")) {
+    param.merge_params.leader_fraction = conf.at("fastener_leader_fraction");
+  }
+  if (conf.contains("fastener_max_leaders")) {
+    param.merge_params.max_leaders = conf.at("fastener_max_leaders");
+  }
+  if (conf.contains("fastener_leaf_size")) {
+    param.merge_params.leaf_size = conf.at("fastener_leaf_size");
+  }
+  if (conf.contains("fastener_leaf_degree")) {
+    param.merge_params.leaf_degree = conf.at("fastener_leaf_degree");
+  }
   param.cagra_params = [conf](raft::matrix_extent<int64_t> extents,
                               cuvs::distance::DistanceType dist_type) {
     // Delayed parsing/initialization of cagra_params - it's called once the dataset shape is known

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.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
@@ -101,6 +101,7 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
     dataset_dependent_params cagra_params;
     size_t num_dataset_splits = 1;
     CagraMergeType merge_type = CagraMergeType::kPhysical;
+    cuvs::neighbors::cagra::merge_params merge_params;
   };
 
   cuvs_cagra(Metric metric, int dim, const build_param& param, int concurrent_searches = 1)
@@ -228,21 +229,10 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
       auto sub_dev =
         raft::make_device_matrix_view<const T, int64_t, raft::row_major>(sub_ptr, rows, dim_);
 
-      auto sub_index = cuvs::neighbors::cagra::index<T, IdxT>(handle_, params.metric);
-      if (index_params_.merge_type == CagraMergeType::kPhysical) {
-        if (dataset_is_on_host) {
-          sub_index.update_dataset(handle_, sub_host);
-        } else {
-          sub_index.update_dataset(handle_, sub_dev);
-        }
-      }
-      if (index_params_.merge_type == CagraMergeType::kLogical) {
-        if (dataset_is_on_host) {
-          sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_host);
-        } else {
-          sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_dev);
-        }
-      }
+      // Build every partition before the merge so FASTENER and REBUILD consume identical prepared
+      // inputs. Partition construction remains outside cagra::merge's NVTX range.
+      auto sub_index = dataset_is_on_host ? cuvs::neighbors::cagra::build(handle_, params, sub_host)
+                                          : cuvs::neighbors::cagra::build(handle_, params, sub_dev);
       auto sub_index_shared =
         std::make_shared<cuvs::neighbors::cagra::index<T, IdxT>>(std::move(sub_index));
       sub_indices_.push_back(std::move(sub_index_shared));
@@ -254,8 +244,8 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
         indices.push_back(ptr.get());
       }
 
-      index_ = std::make_shared<cuvs::neighbors::cagra::index<T, IdxT>>(
-        std::move(cuvs::neighbors::cagra::merge(handle_, params, indices)));
+      index_ = std::make_shared<cuvs::neighbors::cagra::index<T, IdxT>>(std::move(
+        cuvs::neighbors::cagra::merge(handle_, params, indices, index_params_.merge_params)));
     }
   }
 }