Improved CAGRA build parameter heuristics (#1448)

Changes to the build parameter heuristics:
  - Move the code from HNSW namespace to CAGRA namespace to avoid depending on HNSW target
  - Add one more variant of the heuristics: allow generating smaller graph to better match the performance of the HNSW-generated graph
  - Implement automatic switch between NN-Descent and IVF-PQ as the graph-build algorithms depending on the dataset size: NN-Descent tends to perform better on smaller-scale datasets

PR also include C and java bindings.
Resolves #1265

Authors:
  - Artem M. Chirkin (https://github.com/achirkin)
  - Lorenzo Dematté (https://github.com/ldematte)

Approvers:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)
  - Tamas Bela Feher (https://github.com/tfeher)
  - MithunR (https://github.com/mythrocks)
  - Robert Maynard (https://github.com/robertmaynard)

URL: #1448

Author: achirkin

c/include/cuvs/neighbors/cagra.h

@@ -37,6 +37,40 @@ enum cuvsCagraGraphBuildAlgo {
   ITERATIVE_CAGRA_SEARCH = 3
 };
 
+/**
+ * @brief A strategy for selecting the graph build parameters based on similar HNSW index
+ * parameters.
+ *
+ * Define how cuvsCagraIndexParamsFromHnswParams should construct a graph to construct a graph
+ * that is to be converted to (used by) a CPU HNSW index.
+ */
+enum cuvsCagraHnswHeuristicType {
+  /**
+   * Create a graph that is very similar to an HNSW graph in
+   * terms of the number of nodes and search performance. Since HNSW produces a variable-degree
+   * graph (2M being the max graph degree) and CAGRA produces a fixed-degree graph, there's always a
+   * difference in the performance of the two.
+   *
+   * This function attempts to produce such a graph that the QPS and recall of the two graphs being
+   * searched by HNSW are close for any search parameter combination. The CAGRA-produced graph tends
+   * to have a "longer tail" on the low recall side (that is being slightly faster and less
+   * precise).
+   *
+   */
+   CUVS_CAGRA_HEURISTIC_SIMILAR_SEARCH_PERFORMANCE = 0,
+  /**
+   * Create a graph that has the same binary size as an HNSW graph with the given parameters
+   * (graph_degree = 2 * M) while trying to match the search performance as closely as possible.
+   *
+   * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce
+   * the same recalls and QPS for the same parameter ef. The graphs are different internally. For
+   * the same ef, the from-CAGRA index likely has a slightly higher recall and slightly lower QPS.
+   * However, the Recall-QPS curves should be similar (i.e. the points are just shifted along the
+   * curve).
+   */
+   CUVS_CAGRA_HEURISTIC_SAME_GRAPH_FOOTPRINT = 1
+};
+
 /** Parameters for VPQ compression. */
 struct cuvsCagraCompressionParams {
   /**
@@ -145,6 +179,29 @@ cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* param
  */
 cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params);
 
+/**
+ * @brief Create CAGRA index parameters similar to an HNSW index
+ *
+ * This factory function creates CAGRA parameters that yield a graph compatible with
+ * an HNSW graph with the given parameters.
+ *
+ * @param[out] params The CAGRA index params to populate
+ * @param[in] n_rows Number of rows in the dataset
+ * @param[in] dim Number of dimensions in the dataset
+ * @param[in] M HNSW index parameter M
+ * @param[in] ef_construction HNSW index parameter ef_construction
+ * @param[in] heuristic Strategy for parameter selection
+ * @param[in] metric Distance metric to use
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t params,
+                                                int64_t n_rows,
+                                                int64_t dim,
+                                                int M,
+                                                int ef_construction,
+                                                enum cuvsCagraHnswHeuristicType heuristic,
+                                                cuvsDistanceType metric);
+
 /**
  * @}
  */

c/src/neighbors/cagra.cpp

@@ -333,6 +333,63 @@ void get_graph_view(cuvsCagraIndex_t index, DLManagedTensor* graph)
   auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, IdxT>*>(index->addr);
   cuvs::core::to_dlpack(index_ptr->graph(), graph);
 }
+
+// Helper function to populate C IVF-PQ params from C++ params
+static void _populate_c_ivf_pq_params(cuvsIvfPqParams* c_ivf_pq,
+                                    const cuvs::neighbors::cagra::graph_build_params::ivf_pq_params& cpp_ivf_pq)
+{
+  // Populate the IVF-PQ build params
+  auto& bp = cpp_ivf_pq.build_params;
+  c_ivf_pq->ivf_pq_build_params->metric = static_cast<cuvsDistanceType>(bp.metric);
+  c_ivf_pq->ivf_pq_build_params->metric_arg = bp.metric_arg;
+  c_ivf_pq->ivf_pq_build_params->add_data_on_build = bp.add_data_on_build;
+  c_ivf_pq->ivf_pq_build_params->n_lists = bp.n_lists;
+  c_ivf_pq->ivf_pq_build_params->kmeans_n_iters = bp.kmeans_n_iters;
+  c_ivf_pq->ivf_pq_build_params->kmeans_trainset_fraction = bp.kmeans_trainset_fraction;
+  c_ivf_pq->ivf_pq_build_params->pq_bits = bp.pq_bits;
+  c_ivf_pq->ivf_pq_build_params->pq_dim = bp.pq_dim;
+  c_ivf_pq->ivf_pq_build_params->codebook_kind = static_cast<codebook_gen>(bp.codebook_kind);
+  c_ivf_pq->ivf_pq_build_params->force_random_rotation = bp.force_random_rotation;
+  c_ivf_pq->ivf_pq_build_params->conservative_memory_allocation = bp.conservative_memory_allocation;
+  c_ivf_pq->ivf_pq_build_params->max_train_points_per_pq_code = bp.max_train_points_per_pq_code;
+
+  // Populate the IVF-PQ search params
+  auto& sp = cpp_ivf_pq.search_params;
+  c_ivf_pq->ivf_pq_search_params->n_probes = sp.n_probes;
+  c_ivf_pq->ivf_pq_search_params->lut_dtype = sp.lut_dtype;
+  c_ivf_pq->ivf_pq_search_params->internal_distance_dtype = sp.internal_distance_dtype;
+  c_ivf_pq->ivf_pq_search_params->preferred_shmem_carveout = sp.preferred_shmem_carveout;
+
+  c_ivf_pq->refinement_rate = cpp_ivf_pq.refinement_rate;
+}
+
+// Helper function to populate C struct from C++ index_params
+static void _populate_cagra_index_params_from_cpp(cuvsCagraIndexParams_t c_params,
+                                                 const cuvs::neighbors::cagra::index_params& cpp_params)
+{
+  c_params->metric = static_cast<cuvsDistanceType>(cpp_params.metric);
+  c_params->intermediate_graph_degree = cpp_params.intermediate_graph_degree;
+  c_params->graph_degree = cpp_params.graph_degree;
+
+  // Set build algo and parameters based on the variant
+  if (std::holds_alternative<cuvs::neighbors::cagra::graph_build_params::nn_descent_params>(
+        cpp_params.graph_build_params)) {
+    c_params->build_algo = NN_DESCENT;
+    auto nn_params =
+      std::get<cuvs::neighbors::cagra::graph_build_params::nn_descent_params>(
+        cpp_params.graph_build_params);
+    c_params->nn_descent_niter = nn_params.max_iterations;
+  } else if (std::holds_alternative<cuvs::neighbors::cagra::graph_build_params::ivf_pq_params>(
+               cpp_params.graph_build_params)) {
+    c_params->build_algo = IVF_PQ;
+    auto ivf_pq_params =
+      std::get<cuvs::neighbors::cagra::graph_build_params::ivf_pq_params>(
+        cpp_params.graph_build_params);
+
+    _populate_c_ivf_pq_params(c_params->graph_build_params, ivf_pq_params);
+  }
+}
+
 }  // namespace
 
 namespace cuvs::neighbors::cagra {
@@ -665,6 +722,24 @@ extern "C" cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionPar
   return cuvs::core::translate_exceptions([=] { delete params; });
 }
 
+extern "C" cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t params,
+                                                           int64_t n_rows,
+                                                           int64_t dim,
+                                                           int M,
+                                                           int ef_construction,
+                                                           enum cuvsCagraHnswHeuristicType heuristic,
+                                                           cuvsDistanceType metric)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto cpp_metric = static_cast<cuvs::distance::DistanceType>((int)metric);
+    auto cpp_heuristic = static_cast<cuvs::neighbors::cagra::hnsw_heuristic_type>((int)heuristic);
+    auto cpp_params = cuvs::neighbors::cagra::index_params::from_hnsw_params(
+      raft::matrix_extent<int64_t>(n_rows, dim), M, ef_construction, cpp_heuristic, cpp_metric);
+
+    _populate_cagra_index_params_from_cpp(params, cpp_params);
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
 {
   return cuvs::core::translate_exceptions(

cpp/CMakeLists.txt

@@ -409,6 +409,7 @@ if(NOT BUILD_CPU_ONLY)
     src/neighbors/ball_cover/detail/ball_cover/registers_pass_two.cu
     src/neighbors/brute_force.cu
     src/neighbors/brute_force_serialize.cu
+    src/neighbors/cagra.cpp
     src/neighbors/cagra_build_float.cu
     src/neighbors/cagra_build_half.cu
     src/neighbors/cagra_build_int8.cu

cpp/include/cuvs/neighbors/cagra.hpp

@@ -32,6 +32,40 @@ namespace graph_build_params = cuvs::neighbors::graph_build_params;
  * @{
  */
 
+/**
+ * @brief A strategy for selecting the graph build parameters based on similar HNSW index
+ * parameters.
+ *
+ * Define how `cagra::index_params::from_hnsw_params` should construct a graph to construct a graph
+ * that is to be converted to (used by) a CPU HNSW index.
+ */
+enum class hnsw_heuristic_type : uint32_t {
+  /**
+   * Create a graph that is very similar to an HNSW graph in
+   * terms of the number of nodes and search performance. Since HNSW produces a variable-degree
+   * graph (2M being the max graph degree) and CAGRA produces a fixed-degree graph, there's always a
+   * difference in the performance of the two.
+   *
+   * This function attempts to produce such a graph that the QPS and recall of the two graphs being
+   * searched by HNSW are close for any search parameter combination. The CAGRA-produced graph tends
+   * to have a "longer tail" on the low recall side (that is being slightly faster and less
+   * precise).
+   *
+   */
+  SIMILAR_SEARCH_PERFORMANCE = 0,
+  /**
+   * Create a graph that has the same binary size as an HNSW graph with the given parameters
+   * (`graph_degree = 2 * M`) while trying to match the search performance as closely as possible.
+   *
+   * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce
+   * the same recalls and QPS for the same parameter `ef`. The graphs are different internally. For
+   * the same `ef`, the from-CAGRA index likely has a slightly higher recall and slightly lower QPS.
+   * However, the Recall-QPS curves should be similar (i.e. the points are just shifted along the
+   * curve).
+   */
+  SAME_GRAPH_FOOTPRINT = 1
+};
+
 struct index_params : cuvs::neighbors::index_params {
   /** Degree of input graph for pruning. */
   size_t intermediate_graph_degree = 128;
@@ -105,6 +139,39 @@ struct index_params : cuvs::neighbors::index_params {
    * @endcode
    */
   bool attach_dataset_on_build = true;
+
+  /**
+   * @brief Create a CAGRA index parameters compatible with HNSW index
+   *
+   * @param dataset The shape of the input dataset
+   * @param M HNSW index parameter M
+   * @param ef_construction HNSW index parameter ef_construction
+   * @param heuristic The heuristic to use for selecting the graph build parameters
+   * @param metric The distance metric to search
+   *
+   * * IMPORTANT NOTE *
+   *
+   * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce
+   * exactly the same recalls and QPS for the same parameter `ef`. The graphs are different
+   * internally. Depending on the selected heuristics, the CAGRA-produced graph's QPS-Recall curve
+   * may be shifted along the curve right or left. See the heuristics descriptions for more details.
+   *
+   * Usage example:
+   * @code{.cpp}
+   *   using namespace cuvs::neighbors;
+   *   raft::resources res;
+   *   auto dataset = raft::make_device_matrix<float, int64_t>(res, N, D);
+   *   auto cagra_params = cagra::index_params::from_hnsw_params(dataset.extents(), M, efc);
+   *   auto cagra_index = cagra::build(res, cagra_params, dataset);
+   *   auto hnsw_index = hnsw::from_cagra(res, hnsw_params, cagra_index);
+   * @endcode
+   */
+  static cagra::index_params from_hnsw_params(
+    raft::matrix_extent<int64_t> dataset,
+    int M,
+    int ef_construction,
+    hnsw_heuristic_type heuristic       = hnsw_heuristic_type::SIMILAR_SEARCH_PERFORMANCE,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded);
 };
 
 /**

cpp/include/cuvs/neighbors/hnsw.hpp

@@ -57,7 +57,7 @@ struct index_params : cuvs::neighbors::index_params {
  * @brief Create a CAGRA index parameters compatible with HNSW index
  *
  * @param dataset The shape of the input dataset.
- * @param M HNSW index parameter M.
+ * @param M HNSW index parameter M (graph degree = 2*M).
  * @param ef_construction HNSW index parameter ef_construction.
  * @param metric The distance metric to search.
  *
@@ -80,6 +80,7 @@ struct index_params : cuvs::neighbors::index_params {
  *   auto hnsw_index = hnsw::from_cagra(res, hnsw_params, cagra_index);
  * @endcode
  */
+[[deprecated("Use cagra::index_params::from_hnsw_params instead")]]
 cuvs::neighbors::cagra::index_params to_cagra_params(
   raft::matrix_extent<int64_t> dataset,
   int M,

cpp/src/neighbors/cagra.cpp

@@ -0,0 +1,57 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cuvs/distance/distance.hpp>
+#include <cuvs/neighbors/cagra.hpp>
+
+#include <cuvs/neighbors/common.hpp>
+
+namespace cuvs::neighbors::cagra {
+
+inline auto graph_params_heuristic(raft::matrix_extent<int64_t> dataset,
+                                   int intermediate_graph_degree,
+                                   int ef_construction,
+                                   cuvs::distance::DistanceType metric)
+  -> decltype(index_params::graph_build_params)
+{
+  if (dataset.extent(0) < int64_t(1e6)) {
+    // Use NN descent for smaller datasets
+    auto nn_descent_params =
+      graph_build_params::nn_descent_params(intermediate_graph_degree, metric);
+    nn_descent_params.max_iterations = 5 + ef_construction / 16;
+    return nn_descent_params;
+  } else {
+    // Otherwise, use IVF-PQ
+    auto ivf_pq_params = cuvs::neighbors::graph_build_params::ivf_pq_params(dataset, metric);
+    ivf_pq_params.search_params.n_probes =
+      std::round(2 + std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + ef_construction / 16);
+    return ivf_pq_params;
+  }
+}
+
+cagra::index_params index_params::from_hnsw_params(raft::matrix_extent<int64_t> dataset,
+                                                   int M,
+                                                   int ef_construction,
+                                                   hnsw_heuristic_type heuristic,
+                                                   cuvs::distance::DistanceType metric)
+{
+  cagra::index_params params;
+  switch (heuristic) {
+    case hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT:
+      params.graph_degree              = M * 2;
+      params.intermediate_graph_degree = M * 3;
+      break;
+    case hnsw_heuristic_type::SIMILAR_SEARCH_PERFORMANCE:
+    default:
+      params.graph_degree              = 2 + M * 2 / 3;
+      params.intermediate_graph_degree = M + M * ef_construction / 256;
+      break;
+  }
+  params.graph_build_params =
+    graph_params_heuristic(dataset, params.intermediate_graph_degree, ef_construction, metric);
+  return params;
+}
+
+}  // namespace cuvs::neighbors::cagra

cpp/src/neighbors/hnsw.cpp

@@ -4,8 +4,11 @@
  */
 
 #include "detail/hnsw.hpp"
-#include <cstdint>
+
+#include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/hnsw.hpp>
+
+#include <cstdint>
 #include <sys/types.h>
 
 namespace cuvs::neighbors::hnsw {
@@ -15,16 +18,12 @@ auto to_cagra_params(raft::matrix_extent<int64_t> dataset,
                      int ef_construction,
                      cuvs::distance::DistanceType metric) -> cuvs::neighbors::cagra::index_params
 {
-  auto ivf_pq_params = cuvs::neighbors::graph_build_params::ivf_pq_params(dataset, metric);
-  ivf_pq_params.search_params.n_probes =
-    std::round(std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + ef_construction / 16);
-
-  cagra::index_params params;
-  params.graph_build_params        = ivf_pq_params;
-  params.graph_degree              = M * 2;
-  params.intermediate_graph_degree = M * 3;
-
-  return params;
+  return cuvs::neighbors::cagra::index_params::from_hnsw_params(
+    dataset,
+    M,
+    ef_construction,
+    cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT,
+    metric);
 }
 
 #define CUVS_INST_HNSW_FROM_CAGRA(T)                                                  \