Skip to content

Add CAGRA build heuristics independent of HNSW parameters #2345

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 6 commits into from
Jul 28, 2026
Merged

Add CAGRA build heuristics independent of HNSW parameters #2345

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
89d37f5
Add the new heuristics on the C++ side
achirkin Jul 22, 2026
7c78ad3
Add C, Java, and Rust bindings
achirkin Jul 22, 2026
7998bfe
Address docstring suggestions
achirkin Jul 22, 2026
f957708
Merge branch 'release/26.08' into fea-cagra-default-heuristics
achirkin Jul 22, 2026
ba76b0c
Merge branch 'release/26.08' into fea-cagra-default-heuristics
achirkin Jul 23, 2026
c082194
Merge branch 'release/26.08' into fea-cagra-default-heuristics
achirkin Jul 27, 2026
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
23 changes: 22 additions & 1 deletion c/include/cuvs/neighbors/cagra.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -299,6 +299,27 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_
enum cuvsCagraHnswHeuristicType heuristic,
cuvsDistanceType metric);

/**
* @brief Create CAGRA index parameters heuristically tuned for a dataset
*
* This factory function selects the graph build algorithm and its parameters based on the shape of
* the dataset.
*
* @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] graph_degree Degree of the output graph
* @param[in] metric Distance metric to use
* @param[in] build_quality Higher values increase build quality (and cost) up to a point
* @return cuvsError_t
*/
CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params,
int64_t n_rows,
int64_t dim,
size_t graph_degree,
cuvsDistanceType metric,
size_t build_quality);

/**
* @}
*/
Expand Down
18 changes: 17 additions & 1 deletion c/src/neighbors/cagra.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -834,6 +834,22 @@ extern "C" cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t
});
}

extern "C" cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params,
int64_t n_rows,
int64_t dim,
size_t graph_degree,
cuvsDistanceType metric,
size_t build_quality)
{
return cuvs::core::translate_exceptions([=] {
auto cpp_metric = static_cast<cuvs::distance::DistanceType>((int)metric);
auto cpp_params = cuvs::neighbors::cagra::index_params::from_dataset(
raft::matrix_extent<int64_t>(n_rows, dim), graph_degree, cpp_metric, build_quality);

_populate_cagra_index_params_from_cpp(params, cpp_params);
});
}

extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
{
return cuvs::core::translate_exceptions(
Expand Down
62 changes: 56 additions & 6 deletions cpp/include/cuvs/neighbors/cagra.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,12 @@ namespace graph_build_params = cuvs::neighbors::graph_build_params;
* @{
*/

using graph_build_params_t = std::variant<std::monostate,
graph_build_params::ivf_pq_params,
graph_build_params::nn_descent_params,
graph_build_params::ace_params,
graph_build_params::iterative_search_params>;

/**
* @brief A strategy for selecting the graph build parameters based on similar HNSW index
* parameters.
Expand Down Expand Up @@ -181,12 +187,7 @@ struct index_params : cuvs::neighbors::index_params {
* cagra::graph_build_params::iterative_search_params();
* @endcode
*/
std::variant<std::monostate,
graph_build_params::ivf_pq_params,
graph_build_params::nn_descent_params,
graph_build_params::ace_params,
graph_build_params::iterative_search_params>
graph_build_params;
graph_build_params_t graph_build_params;
/**
* Whether to use MST optimization to guarantee graph connectivity.
*/
Expand Down Expand Up @@ -223,6 +224,55 @@ struct index_params : cuvs::neighbors::index_params {
*/
bool attach_dataset_on_build = true;

/**
* @brief Select the graph build algorithm and its parameters for a dataset.
*
* This is the main CAGRA build heuristic: it chooses between NN-descent and IVF-PQ based on the
* dataset size and tunes their parameters based on the target intermediate graph degree and the
* requested build quality. It returns the `graph_build_params` variant only; the caller is
* responsible for setting `graph_degree` / `intermediate_graph_degree`.
*
* @param dataset The shape of the input dataset
* @param intermediate_graph_degree The intermediate (kNN) graph degree the build should target.
* Note: the intermediate graph degree must be not smaller than the output graph degree; a good
* practice is to have it 1.5x to 2x of the desired graph_degree and a multiple of 32.
* @param metric The distance metric to search
* @param build_quality Higher values increase the build quality (and cost) up to a point.
Comment thread
achirkin marked this conversation as resolved.
* Any value is valid, but values below 20 are the most practical (default = 7).
*/
static graph_build_params_t graph_build_heuristic(
raft::matrix_extent<int64_t> dataset,
size_t intermediate_graph_degree,
cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
size_t build_quality = 7);

/**
* @brief Create CAGRA index parameters heuristically tuned for a dataset.
*
* Returns default CAGRA `index_params` with `graph_build_params` selected by
* `graph_build_heuristic` for the given dataset.
*
* @param dataset The shape of the input dataset
* @param graph_degree Degree of the output graph.
* @param metric The distance metric to search
* @param build_quality Higher values increase the build quality (and cost) up to a point.
Comment thread
achirkin marked this conversation as resolved.
* Any value is valid, but values below 20 are the most practical (default = 7).
*
* 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_dataset(dataset.extents());
* auto cagra_index = cagra::build(res, cagra_params, dataset);
* @endcode
*/
static cagra::index_params from_dataset(
raft::matrix_extent<int64_t> dataset,
size_t graph_degree = 64,
cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
size_t build_quality = 7);

/**
* @brief Create a CAGRA index parameters compatible with HNSW index
*
Expand Down
31 changes: 22 additions & 9 deletions cpp/src/neighbors/cagra.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand All @@ -10,27 +10,40 @@

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)
graph_build_params_t index_params::graph_build_heuristic(raft::matrix_extent<int64_t> dataset,
size_t intermediate_graph_degree,
cuvs::distance::DistanceType metric,
size_t build_quality)
{
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;
nn_descent_params.max_iterations = 5 + build_quality;
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);
std::round(2 + std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + build_quality);
return ivf_pq_params;
}
}

cagra::index_params index_params::from_dataset(raft::matrix_extent<int64_t> dataset,
size_t graph_degree,
cuvs::distance::DistanceType metric,
size_t build_quality)
{
cagra::index_params params;
params.metric = metric;
params.graph_degree = graph_degree;
params.intermediate_graph_degree = graph_degree * 3 / 2;
params.graph_build_params =
graph_build_heuristic(dataset, params.intermediate_graph_degree, metric, build_quality);
return params;
}

cagra::index_params index_params::from_hnsw_params(raft::matrix_extent<int64_t> dataset,
int M,
int ef_construction,
Expand All @@ -50,7 +63,7 @@ cagra::index_params index_params::from_hnsw_params(raft::matrix_extent<int64_t>
break;
}
params.graph_build_params =
graph_params_heuristic(dataset, params.intermediate_graph_degree, ef_construction, metric);
graph_build_heuristic(dataset, params.intermediate_graph_degree, metric, ef_construction / 16);
return params;
}

Expand Down
12 changes: 4 additions & 8 deletions cpp/src/neighbors/detail/cagra/cagra_build.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1383,15 +1383,11 @@ index<T, IdxT> build_ace(raft::resources const& res,
std::chrono::duration_cast<std::chrono::milliseconds>(read_end - start).count();

// Create index for this partition
cuvs::neighbors::cagra::index_params sub_index_params;
sub_index_params = cuvs::neighbors::cagra::index_params::from_hnsw_params(
auto sub_index_params = cuvs::neighbors::cagra::index_params::from_dataset(
raft::make_extents<int64_t>(sub_dataset_size, dataset_dim),
graph_degree / 2,
ef_construction,
cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT,
params.metric);
// avoid rounding down graph_degree via heuristics doing `(graph_degree / 2) * 2`
sub_index_params.graph_degree = graph_degree;
graph_degree,
params.metric,
ef_construction / 16);
sub_index_params.attach_dataset_on_build = false;
sub_index_params.guarantee_connectivity = params.guarantee_connectivity;

Expand Down
8 changes: 7 additions & 1 deletion java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package com.nvidia.cuvs;
Expand Down Expand Up @@ -359,6 +359,12 @@ public static CagraIndexParams fromHnswParams(
.cagraIndexParamsFromHnswParams(rows, dim, M, efConstruction, heuristic, metric);
}

public static CagraIndexParams fromDataset(
long rows, long dim, long graphDegree, CuvsDistanceType metric, long buildQuality) {
return CuVSProvider.provider()
.cagraIndexParamsFromDataset(rows, dim, graphDegree, metric, buildQuality);
}

/**
* Gets the degree of input graph for pruning.
*
Expand Down
17 changes: 17 additions & 0 deletions java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,4 +246,21 @@ CagraIndexParams cagraIndexParamsFromHnswParams(
int efConstruction,
CagraIndexParams.HnswHeuristicType heuristic,
CagraIndexParams.CuvsDistanceType metric);

/**
* Create CAGRA index parameters heuristically tuned for a dataset.
*
* @param rows The number of rows in the input dataset
* @param dim The number of dimensions in the input dataset
* @param graphDegree Degree of the output graph
* @param metric The distance metric to search
* @param buildQuality Higher values increase build quality (and cost) up to a point
* @return A new CAGRA index parameters object
*/
CagraIndexParams cagraIndexParamsFromDataset(
long rows,
long dim,
long graphDegree,
CagraIndexParams.CuvsDistanceType metric,
long buildQuality);
}
10 changes: 10 additions & 0 deletions java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,16 @@ public CagraIndexParams cagraIndexParamsFromHnswParams(
throw new UnsupportedOperationException(reasons);
}

@Override
public CagraIndexParams cagraIndexParamsFromDataset(
long rows,
long dim,
long graphDegree,
CagraIndexParams.CuvsDistanceType metric,
long buildQuality) {
throw new UnsupportedOperationException(reasons);
}

@Override
public void setLogLevel(Level level) {
throw new UnsupportedOperationException(reasons);
Expand Down
34 changes: 34 additions & 0 deletions java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,40 @@ public CagraIndexParams cagraIndexParamsFromHnswParams(
}
}

@Override
public CagraIndexParams cagraIndexParamsFromDataset(
long rows,
long dim,
long graphDegree,
CagraIndexParams.CuvsDistanceType metric,
long buildQuality) {
try (var nativeCagraIndexParams = createCagraIndexParams();
var ivfPqIndexParams = createIvfPqIndexParams();
var ivfPqSearchParams = createIvfPqSearchParams()) {

// This is already allocated by cuvsCagraIndexParamsCreate,
// we just need to populate it.
MemorySegment cuvsIvfPqParamsMemorySegment =
cuvsCagraIndexParams.graph_build_params(nativeCagraIndexParams.handle());
cuvsIvfPqParams.ivf_pq_build_params(cuvsIvfPqParamsMemorySegment, ivfPqIndexParams.handle());
cuvsIvfPqParams.ivf_pq_search_params(
cuvsIvfPqParamsMemorySegment, ivfPqSearchParams.handle());

cuvsCagraIndexParams.graph_build_params(
nativeCagraIndexParams.handle(), cuvsIvfPqParamsMemorySegment);
checkCuVSError(
cuvsCagraIndexParamsFromDataset(
nativeCagraIndexParams.handle(), rows, dim, graphDegree, metric.value, buildQuality),
"cuvsCagraIndexParamsFromDataset");

return populateCagraIndexParamsFromNative(
nativeCagraIndexParams,
ivfPqIndexParams,
ivfPqSearchParams,
cuvsIvfPqParamsMemorySegment);
}
}

@Override
public void setLogLevel(Level level) {
if (level.equals(Level.ALL) || level.equals(Level.FINEST)) {
Expand Down
12 changes: 12 additions & 0 deletions rust/cuvs-sys/src/bindings.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1323,6 +1323,18 @@ unsafe extern "C" {
metric: cuvsDistanceType,
) -> cuvsError_t;
}
unsafe extern "C" {
#[must_use]
#[doc = " @brief Create CAGRA index parameters heuristically tuned for a dataset\n\n This factory function selects the graph build algorithm and its parameters based on the shape of\n the dataset.\n\n @param[out] params The CAGRA index params to populate\n @param[in] n_rows Number of rows in the dataset\n @param[in] dim Number of dimensions in the dataset\n @param[in] graph_degree Degree of the output graph\n @param[in] metric Distance metric to use\n @param[in] build_quality Higher values increase build quality (and cost) up to a point\n @return cuvsError_t"]
pub fn cuvsCagraIndexParamsFromDataset(
params: cuvsCagraIndexParams_t,
n_rows: i64,
dim: i64,
graph_degree: usize,
metric: cuvsDistanceType,
build_quality: usize,
) -> cuvsError_t;
}
#[doc = " @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search\n @{\n/\n/**\n @brief Supplemental parameters to extend CAGRA Index\n"]
#[repr(C)]
#[derive(Debug, Copy, Clone)]
Expand Down
Loading