Fast CAGRA Index Merge (#2352)

This PR implements the Fastener graph merge operation. This PR supports merging for `float`, `half`, `int8`, and `uint8` dtypes, and euclidean distances.

There was originally specialization to use int8 GEMM for the integer types, but I ran into portability issues on Ada and it turns out that using the same f32 path for both is simpler and not substantially slower. Currently investigating switching the unified path to TF32 to use tensor cores.

The core logic resides in `cagra_merge_scaffold.cuh`. 

This PR adds 7.64 MiB to `libcuvs.so`, a 2.94% increase.

<img width="2012" height="866" alt="image" src="https://github.com/user-attachments/assets/06249176-042b-42f2-b5ff-366ac316dd1e" />
(H100)
<img width="2310" height="1408" alt="image" src="https://github.com/user-attachments/assets/1b801196-a451-4c77-beab-862c6f1222f7" />
<img width="1024" height="640" alt="image" src="https://github.com/user-attachments/assets/6f320c19-c47b-4564-8781-8872674274bc" />
<img width="1024" height="640" alt="image" src="https://github.com/user-attachments/assets/e197b2a5-4534-403b-9682-5098fcef9c70" />
<img width="1024" height="640" alt="image" src="https://github.com/user-attachments/assets/7818f9d4-174b-4646-8dfc-8529a21e2691" />
<img width="1024" height="640" alt="image" src="https://github.com/user-attachments/assets/43971772-8a1a-4c65-9bcb-606d4c5bb21b" />

Authors:
  - Ben Landrum (https://github.com/landrumb)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)
  - Artem M. Chirkin (https://github.com/achirkin)

URL: #2352

Author: landrumb

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,34 @@ 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] merged_dataset Empty owning dataset handle. Merge first attempts to allocate and
+ *                            populate device storage with the same layout as the input indices. For
+ *                            an unfiltered merge, AUTO and REBUILD can fall back to host storage if
+ *                            device allocation fails; explicit FASTENER reports the allocation
+ *                            failure instead. Keep this dataset alive while using `output_index`.
+ *                            A host-backed output index must be updated with
+ *                            `cuvsCagraUpdateDataset` before device search.
+ * @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/CMakeLists.txt

@@ -1389,6 +1389,7 @@ if(NOT BUILD_CPU_ONLY)
     ${cagra_extend_inst_files}
     src/neighbors/cagra_optimize.cu
     src/neighbors/detail/cagra/graph_shared.cu
+    src/neighbors/detail/cagra/cagra_merge_scaffold_shared.cu
     ${cagra_serialize_inst_files}
     ${cagra_merge_inst_files}
     ${iface_cagra_inst_files}

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

@@ -302,6 +302,18 @@ inline void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::vpq_p
   }
 }
 
+inline void parse_build_param(const nlohmann::json& conf,
+                              cuvs::neighbors::cagra::merge_params& param)
+{
+  param.levels          = conf.value("levels", param.levels);
+  param.root_fanout     = conf.value("root_fanout", param.root_fanout);
+  param.lower_fanout    = conf.value("lower_fanout", param.lower_fanout);
+  param.leader_fraction = conf.value("leader_fraction", param.leader_fraction);
+  param.max_leaders     = conf.value("max_leaders", param.max_leaders);
+  param.leaf_size       = conf.value("leaf_size", param.leaf_size);
+  param.leaf_degree     = conf.value("leaf_degree", param.leaf_degree);
+}
+
 nlohmann::json collect_conf_with_prefix(const nlohmann::json& conf,
                                         const std::string& prefix,
                                         bool remove_prefix = true)
@@ -416,6 +428,20 @@ 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");
+    }
+  }
+  nlohmann::json fastener_conf = collect_conf_with_prefix(conf, "fastener_");
+  if (!fastener_conf.empty()) { parse_build_param(fastener_conf, param.merge_params); }
 
   nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_");
   if (!comp_search_conf.empty()) {

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

@@ -168,6 +168,7 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
     std::optional<cuvs::neighbors::vpq_params> compression = std::nullopt;
     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)
@@ -332,13 +333,14 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
 
       auto sub_index = index_type(handle_, params.metric);
       if (index_params_.merge_type == CagraMergeType::kPhysical) {
-        // Physical merge only needs the rows of every split; cagra::merge builds the graph.
+        // Fastener reuses the input graphs. Build every split so AUTO, FASTENER, and REBUILD all
+        // receive the same prepared indexes.
         if (dataset_is_on_host) {
-          sub_index.update_device_dataset_same_layout(
-            handle_, detail::make_padded_view<T>(handle_, sub_host, sub_dataset_buffer));
+          sub_index = cuvs::neighbors::cagra::build(
+            handle_, params, detail::make_padded_view<T>(handle_, sub_host, sub_dataset_buffer));
         } else {
-          sub_index.update_device_dataset_same_layout(
-            handle_, detail::make_padded_view<T>(handle_, sub_dev, sub_dataset_buffer));
+          sub_index = cuvs::neighbors::cagra::build(
+            handle_, params, detail::make_padded_view<T>(handle_, sub_dev, sub_dataset_buffer));
         }
       }
       if (index_params_.merge_type == CagraMergeType::kLogical) {
@@ -373,8 +375,13 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
       *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));
       // The merged index holds all the rows now; drop the splits rather than keep a second copy
       // of the dataset on the device for the rest of the run.
       sub_indices_.clear();