Skip to content
Merged
Show file tree
Hide file tree
Changes from 15 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
d93331c
Add the new heuristics as static factory functions of cagra::index_pa…
achirkin Oct 22, 2025
e758a0c
Mark new implementations inline to avoid duplicate symbols
achirkin Oct 22, 2025
e873d3a
Add the new functions to the C API
achirkin Oct 22, 2025
b8f9781
Merge branch 'main' into fea-cagra-hnsw-heuristics
achirkin Oct 23, 2025
f15c01a
Merge branch 'main' into fea-cagra-hnsw-heuristics
achirkin Oct 23, 2025
cb06d5a
Fix types in C wrapper code
achirkin Oct 23, 2025
582db6f
CagraIndexParams from C/Cpp heuristics
ldematte Oct 23, 2025
d0ddb1e
Put the new non-template functions in a separate compilation unit
achirkin Oct 23, 2025
8f0605b
Merge branch 'main' into fea-cagra-hnsw-heuristics
achirkin Oct 29, 2025
be3f901
Fix headers
achirkin Oct 29, 2025
5fa89d5
Merge branch 'main' into fea-cagra-hnsw-heuristics
achirkin Oct 30, 2025
ed73a65
Merge branch 'main' into fea-cagra-hnsw-heuristics
achirkin Oct 30, 2025
927d206
Merge remote-tracking branch 'rapidsai/main' into fea-cagra-hnsw-heur…
achirkin Oct 31, 2025
15786b7
Fix bad code merge
achirkin Oct 31, 2025
da8aa38
Make one function out of two
achirkin Oct 31, 2025
3671205
Put C++ helpers in C API implementation in the the unnamed namespace
achirkin Oct 31, 2025
8b4299e
Undo accidental comment reflow
achirkin Nov 3, 2025
9b33fb7
Add documentation to the java wrapper
achirkin Nov 3, 2025
6904106
Prefix the enum value in the C api in lieu of C++ namespaces
achirkin Nov 3, 2025
21dcc9f
Merge branch 'main' into fea-cagra-hnsw-heuristics
achirkin Nov 3, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions c/include/cuvs/neighbors/cagra.h
Original file line number Diff line number Diff line change
Expand Up @@ -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).
*
*/
SIMILAR_SEARCH_PERFORMANCE = 0,
Comment thread
achirkin marked this conversation as resolved.
Outdated
/**
* 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
};

/** Parameters for VPQ compression. */
struct cuvsCagraCompressionParams {
/**
Expand Down Expand Up @@ -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);

/**
* @}
*/
Expand Down
74 changes: 74 additions & 0 deletions c/src/neighbors/cagra.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -665,6 +665,80 @@ extern "C" cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionPar
return cuvs::core::translate_exceptions([=] { delete params; });
}

// Helper function to populate C IVF-PQ params from C++ params
static void populate_c_ivf_pq_params(cuvsIvfPqParams* c_ivf_pq,
Comment thread
achirkin marked this conversation as resolved.
Outdated
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);
}
}

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(
Expand Down
1 change: 1 addition & 0 deletions cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,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
Expand Down
67 changes: 67 additions & 0 deletions cpp/include/cuvs/neighbors/cagra.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
};

/**
Expand Down
3 changes: 2 additions & 1 deletion cpp/include/cuvs/neighbors/hnsw.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand All @@ -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,
Expand Down
57 changes: 57 additions & 0 deletions cpp/src/neighbors/cagra.cpp
Original file line number Diff line number Diff line change
@@ -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
21 changes: 10 additions & 11 deletions cpp/src/neighbors/hnsw.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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) \
Expand Down
Loading