Fast CAGRA Index Merge

c/include/cuvs/neighbors/cagra.h

@@ -223,6 +223,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
  *
@@ -239,6 +260,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
  *
@@ -967,7 +994,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.
@@ -1013,7 +1040,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
@@ -1034,6 +1061,27 @@ CUVS_EXPORT cuvsError_t cuvsCagraMerge(cuvsResources_t res,
                            cuvsDataset_t merged_dataset,
                            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,
+                                                 cuvsDataset_t merged_dataset,
+                                                 cuvsCagraIndex_t output_index);
+
 /**
  * @}
  */

c/src/neighbors/cagra.cpp

@@ -122,6 +122,7 @@ static void merge_indices_for_layout(
   cuvs::neighbors::cagra::index_params const& params_cpp,
   std::vector<cuvs::neighbors::cagra::index<T, uint32_t, DatasetViewT>*>& index_ptrs,
   cuvsFilter filter,
+  cuvs::neighbors::cagra::merge_params const& merge_params,
   cuvsDataset_t merged_dataset,
   cuvsCagraIndex_t output_index)
 {
@@ -151,7 +152,8 @@ static void merge_indices_for_layout(
       auto owner    = std::make_unique<owner_t>(std::move(matrix), dim);
       auto view     = owner->as_dataset_view();
       auto merged_idx =
-        cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs, view, row_filter);
+        cuvs::neighbors::cagra::merge(
+          *res_ptr, params_cpp, index_ptrs, view, merge_params, row_filter);
       auto* holder =
         new cuvs_cagra_c_api_index_lifetime_holder<T, DatasetViewT>{std::move(merged_idx)};
       bind_index_lifetime_holder_to_C_index<T, DatasetViewT>(
@@ -164,7 +166,10 @@ static void merge_indices_for_layout(
       merged_dataset->layout       = output_layout;
       merged_dataset->is_owning    = true;
       return;
-    } catch (std::bad_alloc const&) {
+    } catch (std::bad_alloc const& failure) {
+      if (merge_params.algo == cuvs::neighbors::cagra::merge_algo::FASTENER) {
+        RAFT_FAIL("FASTENER cagra::merge could not allocate device memory: %s", failure.what());
+      }
       // Filtered merge gathers rows with device-only primitives, matching the restriction on the
       // legacy host fallback.
       RAFT_EXPECTS(filter.type == NO_FILTER,
@@ -1132,6 +1137,7 @@ void _merge(cuvsResources_t res,
             cuvsCagraIndex_t* indices,
             size_t num_indices,
             cuvsFilter filter,
+            const cuvs::neighbors::cagra::merge_params& merge_params,
             cuvsDataset_t merged_dataset,
             cuvsCagraIndex_t output_index)
 {
@@ -1179,13 +1185,13 @@ void _merge(cuvsResources_t res,
       convert_opaque_indices_to_concrete_types<T, cuvs::neighbors::device_padded_dataset_view<T, int64_t>>(
         indices, num_indices);
     merge_indices_for_layout<T, cuvs::neighbors::device_padded_dataset_view<T, int64_t>>(
-      res_ptr, params_cpp, index_ptrs, filter, merged_dataset, output_index);
+      res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, output_index);
   } else {
     auto index_ptrs =
       convert_opaque_indices_to_concrete_types<T, cuvs::neighbors::device_standard_dataset_view<T, int64_t>>(
         indices, num_indices);
     merge_indices_for_layout<T, cuvs::neighbors::device_standard_dataset_view<T, int64_t>>(
-      res_ptr, params_cpp, index_ptrs, filter, merged_dataset, output_index);
+      res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, output_index);
   }
 }
 
@@ -1919,15 +1925,41 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res,
                                       cuvsFilter filter,
                                       cuvsDataset_t merged_dataset,
                                       cuvsCagraIndex_t output_index)
+{
+  return cuvsCagraMergeWithParams(
+    res, params, nullptr, indices, num_indices, filter, merged_dataset, 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,
+                                                cuvsDataset_t merged_dataset,
+                                                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");
     RAFT_EXPECTS(indices[0] != nullptr && indices[0]->addr != 0,
                  "All input indices must be built (non-empty)");
 
-    // 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] != nullptr && indices[i]->addr != 0,
@@ -1943,13 +1975,17 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res,
     output_index->addr = 0;
     // Dispatch based on data type
     if (dtype.code == kDLFloat && dtype.bits == 32) {
-      _merge<float>(res, *params, indices, num_indices, filter, merged_dataset, output_index);
+      _merge<float>(
+        res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index);
     } else if (dtype.code == kDLFloat && dtype.bits == 16) {
-      _merge<half>(res, *params, indices, num_indices, filter, merged_dataset, output_index);
+      _merge<half>(
+        res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index);
     } else if (dtype.code == kDLInt && dtype.bits == 8) {
-      _merge<int8_t>(res, *params, indices, num_indices, filter, merged_dataset, output_index);
+      _merge<int8_t>(
+        res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index);
     } else if (dtype.code == kDLUInt && dtype.bits == 8) {
-      _merge<uint8_t>(res, *params, indices, num_indices, filter, merged_dataset, output_index);
+      _merge<uint8_t>(
+        res, *params, indices, num_indices, filter, merge_params_cpp, merged_dataset, output_index);
     } else {
       RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits);
     }
@@ -1995,6 +2031,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

@@ -866,7 +866,12 @@ TEST(CagraC, BuildMergeSearch)
   cuvsCagraIndex_t index_array[2] = {index_main, index_add};
   cuvsDataset_t merged_dataset;
   ASSERT_EQ(cuvsDatasetCreate(&merged_dataset), CUVS_SUCCESS);
-  ASSERT_EQ(cuvsCagraMerge(res, build_params, index_array, 2, filter, merged_dataset, index_merged),
+  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, merged_dataset, index_merged),
             CUVS_SUCCESS);
   {
     cuvsDatasetMemType_t mem_type{};
@@ -950,14 +955,15 @@ 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);
   cuvsCagraIndexDestroy(index_main);
   cuvsDatasetDestroy(padded_dataset_owner);
   cuvsDatasetDestroy(additional_dataset_view);
   cuvsDatasetDestroy(main_dataset_view);
   cuvsDatasetDestroy(merged_dataset);
-  cuvsCagraIndexParamsDestroy(build_params);
   cuvsResourcesDestroy(res);
 }
 

cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h

@@ -416,6 +416,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

@@ -102,6 +102,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)
@@ -307,33 +308,12 @@ 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, static_cast<int64_t>(rows), static_cast<int64_t>(dim_));
 
+      // Build every partition before the merge so FASTENER and REBUILD consume identical
+      // prepared inputs. Unlike the rebuild path, Fastener reuses the input graphs, so every
+      // partition must be fully built here, not merely have a dataset attached.
       auto sub_index = index_type(handle_, params.metric);
-      if (index_params_.merge_type == CagraMergeType::kPhysical) {
-        if (dataset_is_on_host) {
-          sub_dataset_buffers_->emplace_back(
-            raft::make_device_matrix<T, int64_t>(handle_, rows, dim_));
-          raft::copy(sub_dataset_buffers_->back().data_handle(),
-                     sub_ptr,
-                     static_cast<size_t>(rows) * dim_,
-                     raft::resource::get_cuda_stream(handle_));
-          cuvs::neighbors::device_padded_dataset_view<T, int64_t> dv(
-            raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_);
-          sub_index.update_device_dataset_same_layout(handle_, dv);
-        } else {
-          if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dev)) {
-            auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle_, sub_dev);
-            sub_index.update_device_dataset_same_layout(handle_, pdv);
-          } else {
-            auto padded = cuvs::neighbors::make_device_padded_dataset(handle_, sub_dev);
-            sub_dataset_buffers_->push_back(std::move(padded->data_));
-            cuvs::neighbors::device_padded_dataset_view<T, int64_t> pdv(
-              raft::make_const_mdspan(sub_dataset_buffers_->back().view()), dim_);
-            sub_index.update_device_dataset_same_layout(handle_, pdv);
-          }
-        }
-      }
-      if (index_params_.merge_type == CagraMergeType::kLogical) {
-        if (use_ace_host) {
+      {
+        if (index_params_.merge_type == CagraMergeType::kLogical && use_ace_host) {
           // ACE build is always graph-only; build the graph from a host_padded_dataset_view
           // (required by the new build() API), then upload and attach a device padded copy.
           const uint32_t req_stride_sub =
@@ -417,12 +397,18 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
       for (auto* index : indices) {
         merged_rows += static_cast<int64_t>(index->size());
       }
+      // dataset_ is a long-lived member, so the merged index's view stays valid.
       auto const stride        = static_cast<int64_t>(indices.front()->dataset().stride());
       *dataset_                = raft::make_device_matrix<T, int64_t>(handle_, merged_rows, stride);
       auto merged_dataset_view = cuvs::neighbors::device_padded_dataset_view<T, int64_t>(
         raft::make_const_mdspan(dataset_->view()), static_cast<uint32_t>(dim_));
-      index_ = std::make_shared<index_type>(cuvs::neighbors::cagra::merge(
-        handle_, params, indices, merged_dataset_view, merge_row_filter));
+      index_ =
+        std::make_shared<index_type>(cuvs::neighbors::cagra::merge(handle_,
+                                                                   params,
+                                                                   indices,
+                                                                   merged_dataset_view,
+                                                                   index_params_.merge_params,
+                                                                   merge_row_filter));
     }
   }
 }