From 89d37f56f3d5610e1bea815cfce24de933fb25d7 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 22 Jul 2026 07:35:51 +0200 Subject: [PATCH 1/3] Add the new heuristics on the C++ side --- cpp/include/cuvs/neighbors/cagra.hpp | 58 +++++++++++++++++-- cpp/src/neighbors/cagra.cpp | 31 +++++++--- .../neighbors/detail/cagra/cagra_build.cuh | 12 ++-- 3 files changed, 78 insertions(+), 23 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d1937cba27..bd12420775 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -112,6 +112,12 @@ namespace graph_build_params = cuvs::neighbors::graph_build_params; * @{ */ +using graph_build_params_t = std::variant; + /** * @brief A strategy for selecting the graph build parameters based on similar HNSW index * parameters. @@ -181,12 +187,7 @@ struct index_params : cuvs::neighbors::index_params { * cagra::graph_build_params::iterative_search_params(); * @endcode */ - std::variant - graph_build_params; + graph_build_params_t graph_build_params; /** * Whether to use MST optimization to guarantee graph connectivity. */ @@ -223,6 +224,51 @@ 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 + * @param metric The distance metric to search + * @param build_quality Higher values increase the build quality (and cost) up to a point. + */ + static graph_build_params_t graph_build_heuristic( + raft::matrix_extent 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. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * raft::resources res; + * auto dataset = raft::make_device_matrix(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 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 * diff --git a/cpp/src/neighbors/cagra.cpp b/cpp/src/neighbors/cagra.cpp index 6aa5737e36..1de6d5a669 100644 --- a/cpp/src/neighbors/cagra.cpp +++ b/cpp/src/neighbors/cagra.cpp @@ -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 */ @@ -10,27 +10,40 @@ namespace cuvs::neighbors::cagra { -inline auto graph_params_heuristic(raft::matrix_extent 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 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 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 dataset, int M, int ef_construction, @@ -50,7 +63,7 @@ cagra::index_params index_params::from_hnsw_params(raft::matrix_extent 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; } diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index c31c90ce3a..e8b86675d1 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1383,15 +1383,11 @@ index build_ace(raft::resources const& res, std::chrono::duration_cast(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(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; From 7c78ad3662fc13b26f5f29b8f97b377a66572794 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 22 Jul 2026 08:33:07 +0200 Subject: [PATCH 2/3] Add C, Java, and Rust bindings --- c/include/cuvs/neighbors/cagra.h | 23 ++++++++++++- c/src/neighbors/cagra.cpp | 18 +++++++++- .../com/nvidia/cuvs/CagraIndexParams.java | 8 ++++- .../com/nvidia/cuvs/spi/CuVSProvider.java | 17 ++++++++++ .../nvidia/cuvs/spi/UnsupportedProvider.java | 10 ++++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 34 +++++++++++++++++++ rust/cuvs-sys/src/bindings.rs | 12 +++++++ 7 files changed, 119 insertions(+), 3 deletions(-) diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 22809da37e..887b8e7b4c 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -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 */ @@ -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); + /** * @} */ diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 004b810c78..de7f7a5536 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -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 */ @@ -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((int)metric); + auto cpp_params = cuvs::neighbors::cagra::index_params::from_dataset( + raft::matrix_extent(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( diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java index e185ed9f26..7eab7b5485 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java @@ -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; @@ -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. * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 60a2928aec..496d4c5bae 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -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); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 0c5643d2a3..9514c1637e 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -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); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index c9cff1e272..0c9e3feadd 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -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)) { diff --git a/rust/cuvs-sys/src/bindings.rs b/rust/cuvs-sys/src/bindings.rs index 171af6f422..37f3910889 100644 --- a/rust/cuvs-sys/src/bindings.rs +++ b/rust/cuvs-sys/src/bindings.rs @@ -1319,6 +1319,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)] From 7998bfe293392e15057316fbd0ffd7c5c4b8da6b Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 22 Jul 2026 17:16:16 +0200 Subject: [PATCH 3/3] Address docstring suggestions --- cpp/include/cuvs/neighbors/cagra.hpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index bd12420775..1c1ce534dc 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -233,9 +233,12 @@ struct index_params : cuvs::neighbors::index_params { * 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 + * @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. + * 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 dataset, @@ -253,6 +256,7 @@ struct index_params : cuvs::neighbors::index_params { * @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. + * Any value is valid, but values below 20 are the most practical (default = 7). * * Usage example: * @code{.cpp}