Iterative CAGRA-Q by irina-resh-nvda · Pull Request #1810 · NVIDIA/cuvs

irina-resh-nvda · 2026-02-16T15:08:36Z

Build CAGRA on PQ datasets with Iterative CAGRA-Q

Iterative cagra graph construction using CAGRA-Q search.

This PR improves the iterative CAGRA build method by enabling PQ compression: the dataset is compressed before the iterative search starts, and CAGRA-Q is used to iteratively update the KNN graph.

This is the first time we are introducing building CAGRA on (PQ) quantized datasets directly.

copy-pr-bot · 2026-02-16T15:08:40Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

… search - Configurable growth-phase in-build search params (itopk_size, search_width, max_iterations) and internal/smem dtype; itopk auto-forced on the final full-size iteration. - Decouple compression params used during iterative construction from the target index compression. - Add shuffle_dataset option; fix out-of-bounds access from the in-place raft gather by switching to an out-of-place gather.

…around) The shuffle_dataset path used an out-of-place gather into a temporary buffer to work around an illegal memory access in raft's in-place gather overload when n_rows * row_len exceeded 2^31 (32-bit index overflow). That bug is now fixed upstream in raft (NVIDIA/raft#3059, closes #3055), which the cuvs raft pin now includes. Revert to the in-place gather to drop the extra full-size temporary allocation and copy.

…set on disk. Also added tests

…t to the algorithm as a path

…_dataset() is always true

…ches without the dense rows

achirkin

Please undo all changes to bench/ann/src/common. Let's keep CAGRA-Q compressed dataset support as before - by letting the compressed dataset be a part of index file / folder.
We should leave a proper refactoring of the benchmark harness for discussion and implementation in follow up work (if and how we want to pass compressed datasets as first-class input to the benchmarks).

…e_cagra_q

irina-resh-nvda · 2026-08-18T12:42:32Z

Please undo all changes to bench/ann/src/common. Let's keep CAGRA-Q compressed dataset support as before - by letting the compressed dataset be a part of index file / folder. We should leave a proper refactoring of the benchmark harness for discussion and implementation in follow up work (if and how we want to pass compressed datasets as first-class input to the benchmarks).

I created a standalone PR for the vpq dataset serialisation. We can look through it and merge faster that this PR, and start working on the cuvsbench refactoring sooner.

I need the refactoring to rerun the benchmarks (the older performance numbers are obsolete now). I won't be comfortable merging this PR into main before the new benchmarks land

… graphs are already on the device

…ing the CAGRA-Q path

tarang-jain

Thanks for the great work!

tarang-jain · 2026-09-02T20:51:53Z

+  auto batch_queries =
+    vpq_queries != nullptr || dev_query_view.extent(1) != query_dim
+      ? raft::make_device_matrix<T, int64_t>(res, static_cast<int64_t>(max_chunk_size), query_dim)
+      : raft::make_device_matrix<T, int64_t>(res, 0, 0);


Looks like CAGRA search will pad these queries to the aligned dim if they are not aligned right? So why not make the batch itself padded here -- will avoid an extra D2D copy during cagra search (it might need the reconstruct kernel to jump the padded dim).

tarang-jain · 2026-09-02T21:51:15Z

What do you think about putting the ANN bench changes in a separate PR (even though you want to benchmark this e2e)? That can be a follow up to this PR.

tfeher

Thanks @irina-resh-nvda and @aamijar for this PR! My main questions below are about naming of params / APIs related to PQ compressed input. Otherwise, the PR looks great overall.

tfeher · 2026-09-05T06:23:26Z

+ * @brief Compress a dense dataset into a device VPQ dataset.
+ */
+CUVS_EXPORT cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res,
+                                           cuvsCagraCompressionParams_t params,


This function compresses the dataset into VPQ format. The user can decide later to use that for CAGRA, or for other algorithms. It would make sense to rename the params struct to something like cuvsPqParams.

tfeher · 2026-09-05T06:53:47Z

+        if (build_search_conf.contains("persistent")) {
+          arg.persistent = build_search_conf.at("persistent");
+        }
+        if (build_search_conf.contains("persistent_lifetime")) {
+          arg.persistent_lifetime = build_search_conf.at("persistent_lifetime");
+        }
+        if (build_search_conf.contains("persistent_device_usage")) {
+          arg.persistent_device_usage = build_search_conf.at("persistent_device_usage");


I assume build_search is set of parameters used only during build. Is persistent mode used during build? I don't expect that, since with large batch search operations we have large throughput without the persistent kernel. My concern in exposing this as tunable build parameter is that it might motivate users to add these options to their build parameter scan, which is not productive.

If the goal is completeness in parameter handling, then shouldn't we use the same function that parses search params?

tfeher · 2026-09-07T15:52:59Z

    if (out_dataset == nullptr) {
-      cuvs::neighbors::detail::skip_dense_dataset<T, int64_t>(res, is);
+      // Dropping the rows leaves a searchable index for a dense view, whose dataset can be
+      // reattached from the caller's own copy, but not for a VPQ one: the compressed rows exist


Why do we handle the pq compressed dataset differently? Can't the user reattach a pq compressed dataset to an index after the graph was loaded from file?

tfeher · 2026-09-07T16:28:03Z

+    final_graph_size = static_cast<uint64_t>(dataset.n_rows());
+    vpq_dataset      = &dataset.dset();
+  } else {
+    auto search_dataset = ensure_device_padded_for_iterative_search<T>(res, dataset, padded_own);


I believe that using the new DataSet API, the guideline is that we do not make such copies for the user, but expect them to provide already padded data.

so we basically error out if its not padded, right?

tfeher · 2026-09-07T16:42:57Z

+                                                                         logical_dim);
+      auto idx = cuvs::neighbors::cagra::device_padded_index<T, IdxT>(res, params.metric);
+      idx.update_device_dataset_same_layout(res, sub_padded);
+      if (use_device_graph) {


Expect of the first iteration, we always use dev_graph, right? Could you simply copy the initial cagra_graph to the device so that we can simplify the code here?

Keep VPQ serialize/deserialize and iterative-build search_main usage, and adopt KvikIO dataset I/O plus cuda_rt_essentials from the release branch.

aamijar · 2026-09-09T00:46:11Z

+    index_type idx(res, effective_params.metric);
+    idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view()));
+    if (effective_params.attach_dataset_on_build) {
+      idx.update_device_dataset_same_layout(res, dataset);


I thought we got rid of this function update_device_dataset_same_layout, probably a merge conflict artifact

Those build_search_* knobs are preserved on iterative-cagra-q-ann-bench.

github-project-automation Bot added this to Unstructured Data Processing Feb 16, 2026

aamijar assigned irina-resh-nvda Feb 24, 2026

aamijar added non-breaking Introduces a non-breaking change feature request New feature or request labels Feb 24, 2026

cjnolet moved this to In Progress in Unstructured Data Processing Feb 27, 2026

josephine-wolf-oberholtzer added this to Unstructured Data Processing Jul 1, 2026

josephine-wolf-oberholtzer moved this to In Progress in Unstructured Data Processing Jul 1, 2026

irina-resh-nvda added 7 commits July 14, 2026 00:53

chore(cagra): drop local dev artifacts (.clangd, .gitignore entries)

23aed4e

feat(cagra): add batched_device_view_from_host utility and unit test

a4ebee6

fix(cagra): bound random seed selection to graph size during build

90fcf80

feat(bench): expose iterative CAGRA-Q build/search params in cuvs_bench

90d1219

test(cagra): VPQ/iterative build test updates

a938958

irina-resh-nvda force-pushed the iterative_cagra_q branch from 8fc10ac to 4f4068a Compare July 15, 2026 08:46

Merge branch 'main' into iterative_cagra_q

c804088

irina-resh-nvda marked this pull request as ready for review July 15, 2026 08:48

irina-resh-nvda requested review from a team as code owners July 15, 2026 08:48

aamijar and others added 2 commits July 22, 2026 14:57

Merge branch 'main' into iterative_cagra_q

278a4f3

fix style

cc65291

aamijar changed the title ~~Iterative cagra q~~ Iterative graph build using CAGRA-Q search Jul 22, 2026

aamijar changed the title ~~Iterative graph build using CAGRA-Q search~~ Iterative CAGRA-Q Jul 22, 2026

aamijar added 5 commits July 22, 2026 22:19

restore clangd and gitignore changes

25e6d7d

revert cuvs_bench warning

6bca275

remove whitespace

6ed4ab5

remove duplicate file in cmakelists.txt

9d44b23

revert to auto for type deduction

73c9bbc

aamijar and others added 7 commits August 16, 2026 01:26

remove some small diff

79bd7bb

VPQ dataset serialization so CAGRA-Q can build from a compressed data…

bb1fcde

…set on disk. Also added tests

Recognise a compressed .vpq base_file in the ann benchmark and hand i…

31c6322

…t to the algorithm as a path

Don't refuse a compressed base set in search mode: cuvs_cagra's needs…

843425b

…_dataset() is always true

Serialize the VPQ dataset with a CAGRA-Q index so a loaded index sear…

689a882

…ches without the dense rows

C and Python APIs

8863d1d

Merge branch 'main' into iterative_cagra_q

10c7e1f

achirkin requested changes Aug 18, 2026

View reviewed changes

Merge remote-tracking branch 'origin/iterative_cagra_q' into iterativ…

e47bbca

…e_cagra_q

irina-resh-nvda added 5 commits August 18, 2026 05:44

Reverse cuvsbench changes

28ca46b

Share the iterative CAGRA build query chunk size as a constant

682ee36

Don't count graph staging buffers in optimize_workspace_size when the…

b3c551f

… graphs are already on the device

Model the iterative CAGRA build's peak host and device memory, includ…

c030c57

…ing the CAGRA-Q path

added tests for memory estimation

f4790ca

aamijar self-assigned this Aug 20, 2026

merge main and fix conflicts

7d9db3b

tarang-jain reviewed Sep 2, 2026

View reviewed changes

tfeher requested changes Sep 7, 2026

View reviewed changes

tarang-jain changed the base branch from main to release/26.10 September 8, 2026 21:42

tarang-jain added 5 commits September 8, 2026 15:35

rename

6cc01fd

pre-allocate buffers

c45f76d

try using strided queries

4f5e9eb

Merge NVIDIA release/26.10 into iterative_cagra_q.

66b82ed

Keep VPQ serialize/deserialize and iterative-build search_main usage, and adopt KvikIO dataset I/O plus cuda_rt_essentials from the release branch.

bring graph to device

bdc1b84

aamijar reviewed Sep 9, 2026

View reviewed changes

tarang-jain and others added 3 commits September 8, 2026 17:48

Drop iterative CAGRA-Q ann-bench param parsing from this branch.

ec4bc35

Those build_search_* knobs are preserved on iterative-cagra-q-ann-bench.

more review changes

70180eb

Merge branch 'release/26.10' into iterative_cagra_q

3166185

Uh oh!

Conversation

irina-resh-nvda commented Feb 16, 2026 • edited by aamijar Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Build CAGRA on PQ datasets with Iterative CAGRA-Q

Uh oh!

copy-pr-bot Bot commented Feb 16, 2026

Uh oh!

achirkin left a comment

Choose a reason for hiding this comment

Uh oh!

irina-resh-nvda commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tarang-jain left a comment

Choose a reason for hiding this comment

Uh oh!

tarang-jain Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

tarang-jain commented Sep 2, 2026

Uh oh!

tfeher left a comment

Choose a reason for hiding this comment

Uh oh!

tfeher Sep 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tfeher Sep 5, 2026

Choose a reason for hiding this comment

Uh oh!

tfeher Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

tfeher Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

tarang-jain Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

tfeher Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

aamijar Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

irina-resh-nvda commented Feb 16, 2026 •

edited by aamijar

Loading

irina-resh-nvda commented Aug 18, 2026 •

edited

Loading

aamijar Sep 9, 2026 •

edited

Loading