[WIP] Fast CAGRA Index Merge

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)));
     }
   }
 }

cpp/include/cuvs/neighbors/cagra.hpp

@@ -697,6 +697,16 @@ struct CUVS_EXPORT index : cuvs::neighbors::index {
     graph_view_ = knn_graph;
   }
 
+  /**
+   * Replace the graph by taking ownership of an existing device matrix.
+   */
+  void update_graph(raft::resources const&,
+                    raft::device_matrix<graph_index_type, int64_t, raft::row_major>&& knn_graph)
+  {
+    graph_      = std::move(knn_graph);
+    graph_view_ = graph_.view();
+  }
+
   /**
    * Replace the graph with a new graph.
    *
@@ -2447,42 +2457,65 @@ void serialize_to_hnswlib(
  * @{
  */
 
-/** @brief Merge multiple CAGRA indices into a single index.
- *
- * This function merges multiple CAGRA indices into one, combining both the datasets and graph
- * structures.
- *
- * @note: When device memory is sufficient, the dataset attached to the returned index is allocated
- * in device memory by default; otherwise, host memory is used automatically.
- *
- * @note: This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`), and
- * attempting a logical merge here will throw an error.
- *
- * Usage example:
- * @code{.cpp}
- *   using namespace cuvs::neighbors;
- *   auto dataset0 = raft::make_host_matrix<float, int64_t>(handle, size0, dim);
- *   auto dataset1 = raft::make_host_matrix<float, int64_t>(handle, size1, dim);
+enum class merge_algo {
+  /** Select Fastener when its complete preflight succeeds, otherwise preserve rebuild behavior. */
+  AUTO,
+  /** Require Fastener and reject unsupported configurations. */
+  FASTENER,
+  /** Concatenate datasets and rebuild the graph. */
+  REBUILD,
+};
+
+/** C++ controls for physical CAGRA index merge. */
+struct merge_params {
+  merge_algo algo        = merge_algo::AUTO;
+  uint32_t levels        = 2;
+  uint32_t root_fanout   = 2;
+  uint32_t lower_fanout  = 3;
+  double leader_fraction = 0.02;
+  uint32_t max_leaders   = 1000;
+  uint32_t leaf_size     = 256;
+  uint32_t leaf_degree   = 4;
+};
+
+/** @brief Merge multiple physical CAGRA indices into one.
  *
- *   auto index0 = cagra::build(res, index_params, dataset0);
- *   auto index1 = cagra::build(res, index_params, dataset1);
+ * The overload without `merge_params` uses `merge_algo::AUTO`. AUTO runs Fastener only after a
+ * non-mutating preflight validates every input and option; otherwise it calls the existing rebuild
+ * implementation. REBUILD always preserves every input dataset.
  *
- *   std::vector<cagra::index<float, uint32_t>*> indices{&index0, &index1};
+ * Fastener supports unfiltered, uncompressed `float`, `half`, `int8_t`, and `uint8_t`
+ * indices using L2Expanded and `uint32_t` graph IDs. Fastener uses `root_fanout` at the first
+ * split and `lower_fanout` at every later split. Fanouts from 1 through 32, leader fractions in
+ * (0, 1], leader caps through 8192, leaf sizes from 1 through 256, and leaf degrees from 1
+ * through 8 are supported. The configured spill width times `leaf_degree` must not exceed 255.
+ * `index_params::graph_degree` is the final output degree.
  *
- *   auto merged_index = cagra::merge(res, index_params, indices);
- * @endcode
+ * Fastener copies the input datasets into one contiguous device allocation and never mutates the
+ * input indices. When `attach_dataset_on_build` is true, the result owns a 16-byte-row-aligned
+ * copy of that consolidated dataset (matching `cagra::build` / `update_dataset`).
  *
- * @param[in] res RAFT resources used for the merge operation.
- * @param[in] params Parameters that control the merging process.
- * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must:
- *                    - Have attached datasets with the same dimension.
- * @param[in] row_filter an optional device filter function object that greenlights rows
- *    to include in the merged index  (none_sample_filter for no filtering)
- * @return A new CAGRA index containing the merged indices, graph, and dataset.
+ * @param[in] res RAFT resources used for the merge.
+ * @param[in] params Parameters for the returned CAGRA index.
+ * @param[in] indices CAGRA indices to merge.
+ * @param[in] row_filter Optional row filter. Any filter selects rebuild in AUTO and is rejected by
+ * explicit FASTENER.
+ * @return The merged physical CAGRA index.
+ */
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::index<float, uint32_t>*>& indices,
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::index<float, uint32_t>;
+
+/** @copydoc merge
+ * @param[in] merge_params Parameters for the merge.
  */
 auto merge(raft::resources const& res,
            const cuvs::neighbors::cagra::index_params& params,
            std::vector<cuvs::neighbors::cagra::index<float, uint32_t>*>& indices,
+           const cuvs::neighbors::cagra::merge_params& merge_params,
            const cuvs::neighbors::filtering::base_filter& row_filter =
              cuvs::neighbors::filtering::none_sample_filter{})
   -> cuvs::neighbors::cagra::index<float, uint32_t>;
@@ -2495,6 +2528,15 @@ auto merge(raft::resources const& res,
              cuvs::neighbors::filtering::none_sample_filter{})
   -> cuvs::neighbors::cagra::index<half, uint32_t>;
 
+/** @copydoc merge */
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::index<half, uint32_t>*>& indices,
+           const cuvs::neighbors::cagra::merge_params& merge_params,
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::index<half, uint32_t>;
+
 /** @copydoc merge */
 auto merge(raft::resources const& res,
            const cuvs::neighbors::cagra::index_params& params,
@@ -2503,10 +2545,28 @@ auto merge(raft::resources const& res,
              cuvs::neighbors::filtering::none_sample_filter{})
   -> cuvs::neighbors::cagra::index<int8_t, uint32_t>;
 
+/** @copydoc merge */
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::index<int8_t, uint32_t>*>& indices,
+           const cuvs::neighbors::cagra::merge_params& merge_params,
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::index<int8_t, uint32_t>;
+
+/** @copydoc merge */
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::index<uint8_t, uint32_t>*>& indices,
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::index<uint8_t, uint32_t>;
+
 /** @copydoc merge */
 auto merge(raft::resources const& res,
            const cuvs::neighbors::cagra::index_params& params,
            std::vector<cuvs::neighbors::cagra::index<uint8_t, uint32_t>*>& indices,
+           const cuvs::neighbors::cagra::merge_params& merge_params,
            const cuvs::neighbors::filtering::base_filter& row_filter =
              cuvs::neighbors::filtering::none_sample_filter{})
   -> cuvs::neighbors::cagra::index<uint8_t, uint32_t>;

cpp/src/neighbors/cagra.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
  */
 
@@ -427,6 +427,16 @@ index<T, IdxT> merge(raft::resources const& handle,
   return cagra::detail::merge<T, IdxT>(handle, params, indices, row_filter);
 }
 
+template <class T, class IdxT>
+index<T, IdxT> merge(raft::resources const& handle,
+                     const cagra::index_params& params,
+                     std::vector<cuvs::neighbors::cagra::index<T, IdxT>*>& indices,
+                     const cagra::merge_params& merge_params,
+                     const cuvs::neighbors::filtering::base_filter& row_filter)
+{
+  return cagra::detail::merge<T, IdxT>(handle, params, indices, merge_params, row_filter);
+}
+
 /** @} */  // end group cagra
 
 }  // namespace cuvs::neighbors::cagra
@@ -439,4 +449,14 @@ index<T, IdxT> merge(raft::resources const& handle,
     -> cuvs::neighbors::cagra::index<T, IdxT>                                           \
   {                                                                                     \
     return cuvs::neighbors::cagra::merge<T, IdxT>(handle, params, indices, row_filter); \
+  }                                                                                     \
+  auto merge(raft::resources const& handle,                                             \
+             const cuvs::neighbors::cagra::index_params& params,                        \
+             std::vector<cuvs::neighbors::cagra::index<T, IdxT>*>& indices,             \
+             const cuvs::neighbors::cagra::merge_params& merge_params,                  \
+             const cuvs::neighbors::filtering::base_filter& row_filter)                 \
+    -> cuvs::neighbors::cagra::index<T, IdxT>                                           \
+  {                                                                                     \
+    return cuvs::neighbors::cagra::merge<T, IdxT>(                                      \
+      handle, params, indices, merge_params, row_filter);                               \
   }