From d8a6d892e211244e09cbcf5c409ee75ad8077f8d Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:41:58 +0200 Subject: [PATCH] Add CAGRA build heuristics independent of HNSW parameters (#2345) Split the `cagra::index_params::from_hnsw_params` into a core `cagra::index_params::from_dataset` and a thin wrapper ``cagra::index_params::from_hnsw_params`. This solves two problems: 1. Avoid using the HNSW parameters that do not make sense for CAGRA 2. Allow HNSW parameter heuristics change without affecting CAGRA in future The new parameter defaults are selected to yield exactly the same configurations at old call sites of `from_hnsw_params` Authors: - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Tamas Bela Feher (https://github.com/tfeher) URL: https://github.com/NVIDIA/cuvs/pull/2345 --- c/include/cuvs/neighbors/cagra.h | 23 ++++++- c/src/neighbors/cagra.cpp | 18 +++++- cpp/include/cuvs/neighbors/cagra.hpp | 62 +++++++++++++++++-- cpp/src/neighbors/cagra.cpp | 31 +++++++--- .../neighbors/detail/cagra/cagra_build.cuh | 12 ++-- .../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 ++++ 10 files changed, 201 insertions(+), 26 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/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d1937cba27..1c1ce534dc 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,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. + * 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, + 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. + * 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(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; 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 58b61e949b..4576b2bfb1 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)]