Fix benchmark regressions from #1846 by achirkin · Pull Request #2384 · NVIDIA/cuvs

achirkin · 2026-07-31T10:36:15Z

Fix a few issues introduced to the benchmark executables limiting the functionality:

Reinstate the dropped support for CAGRA-Q (compression_* parameters)
Avoid multiple copies of the data and the graph on device in various code paths
Avoid copying the data if it is device-accessible
Make logical merge reuse one buffer slot per split and free it before reallocating, instead of appending new split copies on every set_search_dataset call.
Make multi-GPU CAGRA benchmarks use the padded index and dataset view, which CAGRA search requires.
Revert copy() to the shallow copy-constructor form used before, so all shared state is forwarded to per-thread wrappers.

The changeset is limited to C++ benchmarks, thus doesn't affect the library or any dependencies.

### Overview Addressing NVIDIA#1574 and NVIDIA#1571. Replaced strided_dataset with padded_dataset class. Added support all the way up to CAGRA code. ### Old class structure (Classes + Inheritance): <img width="973" height="333" alt="Screenshot 2026-06-23 at 7 06 37 PM" src="https://github.com/user-attachments/assets/470405f1-00ae-4d02-a5f2-aa11612d2c91" /> ### New Class Structure (ContainerType Tags + Composition): <img width="2012" height="141" alt="Screenshot 2026-06-23 at 7 07 23 PM" src="https://github.com/user-attachments/assets/45e3a17f-14e7-4cf4-8d38-eadc2a1c04d3" /> Inheritance is removed entirely and all dataset types are on the same level of the inheritance tree. 3 Levels: 1) Storage <img width="862" height="442" alt="Screenshot 2026-07-06 at 11 03 41 AM" src="https://github.com/user-attachments/assets/2f24167f-88a3-4b81-9cd4-ca6ba16d93e4" /> 2) Container tags <img width="839" height="250" alt="Screenshot 2026-07-06 at 11 03 11 AM" src="https://github.com/user-attachments/assets/fcb2bda2-a98b-43ca-946c-fec03ce40aaa" /> 3) Public dataset/view aliases <img width="496" height="39" alt="Screenshot 2026-06-23 at 7 02 34 PM" src="https://github.com/user-attachments/assets/6aaff383-2060-4f4a-bafb-a0f258e81f68" /> ### Ownership The index and cagra::build / cagra::index do not own raw vector storage, they only take views. - callers (or the C merged holder) must keep backing memory alive for as long as the index is used. The old code had a type-erased std::unique_ptr<dataset_view<...>>, i.e. non-owning view handles. The new code uses templates on the index type which determines the type of dataset_view the index holds. ### ACE v.s. non-ACE paths on Host ACE path copies datasets that can't entirely fit in CPU memory in chunks onto GPU memory by calling make_padded_dataset. This is 1x memory on CPU and 1x memory on GPU. ### Return types: Used mainly to maintain lifetime of dataset. cuvs_cagra_c_api_lifetime_holder - unique_ptr<vpq_dataset> vpq_owner - unique_ptr<dataset> padded_dataset_owner - raft::device_matrix dataset - cagra::index idx It is a single C++ struct in cagra.cpp that groups the real cagra::index with any extra heap-owned things the C API had to create so the index’s non-owning views stay valid. ### Miscellaneous: Extend Serialize Deserialize Will fill in later ### Factories: - make_device_padded_dataset_view - make_host_padded_dataset_view - make_device_padded_dataset - make_host_padded_dataset - make_vpq_dataset - in pq.hpp and pq.cu - make_merged_dataset - in cagra.hpp ### Places where make_padded_dataset/view are called internally (not by user): Host non-ACE path - cpp/src/neighbors/cagra_build_inst.cu.in - cagra_from_host_padded in cpp/src/neighbors/iface/iface.hpp - c/src/neighbors/cagra.cpp Tiered CAGRA - update_cagra_ann_dataset_for_stride - build_upstream_ann ### Ownership in Downstream Functions: - build() takes dataset_view only. - Downstream functions search / serialize / deserialize / merge only take views. ### Improvements: - build() function previously only supported device dataset inputs. It now supports host dataset inputs. ### Breaking Changes for Dataset API: The following functions are removed since index no longer owns the dataset, index only takes views: - Removed all owning dataset based builds. Build only takes views. - Removed all update_dataset() overloads that take owning dataset. Update_dataset() only takes views. Removed old functions that took mdspan or derivatives of mdspan. ### 4 cases where index previously owned dataset [all deprecated paths]: 2 edge case build() paths when attach_dataset_on_build == true and a successful dense attach: - Non-ACE / typical padded attach: rows live under index_owning_dataset_storage_ (type-erased owning wrapper, commonly device_padded_dataset). - ACE in-memory device_matrix attach: rows live under index_owning_dataset_storage_ (optional raw device_matrix). Compression Param: - implicit vpq dataset creation within build() when compression==True - this forces index to own new vpq dataset which violates our new contract that we want index to never own dataset. Merge: - MERGE path: merge() internally creates merged_dataset on a deprecated internal merged dataset creation path. Here, index takes ownership of merged_dataset by storing it in index_owning_dataset_storage_ . These paths have since been removed. ### Attach Dataset - Previously, in the old code, ACE attach_dataset_on_build called make_device_padded_dataset on host dataset which did a H2D copy in order to attach dataset to final index. - cpp/src/neighbors/detail/cagra/cagra_build.cuh - This has since been removed. attach_dataset_on_build is disabled for host build paths. This avoids a H2D copy. ### Compressed Dataset - Removed old code that did compressed dataset creation within build. This should only happen in the factory. ### Merged Dataset - Removed implicit memory allocation within merge(), memory allocation now delegated to make_merged_dataset() factory. Removed index ownership within merge. ### Deserialize - Removed index ownership of dataset during deserialization. Now users are expected to create/declare the dataset type to be deserialized and then pass it as a reference to the deserialize() function which will then populate this dataset and return it to the caller. ### Helpers - cagra_required_row_width - matrix_actual_row_width - matrix_row_width_matches_cagra_required - convert_dataset_view_to_padded_for_graph_build - convert_host_to_device_index - attach_device_dataset_on_host_index #### How to attach a compressed dataset onto an uncompressed index? 1. construct a new compressed index 2. copy over the graph and other params from the uncompressed index 3. delete the old uncompressed index 4. attach the vpq dataset onto the compressed index #### How to attach a searchable device dataset onto an index built with host build? 1. Convert host index to device index with helper function convert_host_to_device_index a. Utilizes map of host dataset type to device dataset type counterpart ### TODOs: - Bring back Host functions [DONE] - Mark any old functions that are no longer used as [DONE] - Use templates wherever possible. Shift towards composition rather than inheritance [DONE] ### Recent Updates: - build_ace() and build() functions merged on Public API surface - removed build_result(), ace_build_result(), and merge_result() - deprecated internal vpq dataset creation inside build() when index_params::compression == true --> moved to make_vpq_dataset() factory - deprecated internal merged dataset creation inside merge() --> moved to make_merged_dataset() factory. For backwards compatibility, have index take ownership of deprecated internal merged dataset creation path ONLY. - build() and build_ace() both had a attach_dataset_on_build which requires ownership of dataset. Ownership is given to index temporarily. This will later be deprecated. Users will be expected to pass padded dataset on device and call search() directly. Attach_dataset_on_build will no longer be supported for host builds. - added host versions of dataset API - templated build() and downstream functions to work on host datasets - added index template type conversion helpers ### Future PRs: PR#2: Add Support for Compressed Datasets - pq_dataset - bbq_dataset - rabitq_dataset - sq_dataset PR#3: Migrate Rest of Algorithms to use Dataset API - HNSW - IVF - Vamana - Scann - Brute Force --------- Co-authored-by: aamijar <aamijar230@gmail.com> Co-authored-by: divyegala <divyegala@gmail.com> Co-authored-by: Yan Zaretskiy <yzaretskiy@nvidia.com> Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com> Co-authored-by: tarangj <jaintarang2015@gmail.com> Co-authored-by: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Co-authored-by: Igor Motov <igor@motovs.org>

copy-pr-bot · 2026-07-31T10:36:19Z

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.

achirkin · 2026-07-31T10:36:28Z

/ok to test

achirkin · 2026-07-31T11:30:08Z

/ok to test

cjnolet · 2026-07-31T12:06:17Z

@achirkin its too late for anything to be merged into 26.08. You'll need to retarget to main.

…enchmark-provided dataset to avoid ambiguity w.r.t. data residence

achirkin · 2026-07-31T15:10:51Z

Sure. I'll just finish it up today and let you decide whether to re-target and whether to optionally backport the fixes later if necessary.

tfeher

Thanks Artem, most of it looks good, but I have some questions.

tfeher · 2026-08-03T11:31:55Z

+ * Turn a host-resident index into a searchable device index carrying only the graph.
+ *
+ * The graph is already in device memory, owned by `host_index`; the returned index only views it,
+ * so the caller must keep `host_index` alive. The dataset is attached later, by `set_search_param`
+ * or `set_search_dataset`.


Could you explain why is this function needed? What does it mean "host_index"? An index that stores the graph in host memory? But the description says "The graph is already in device memory"

That means the new cuvs::neighbors::cagra::host_standard_index<T, IdxT> type - so it assumes the dataset is on the host (template parameter), while the graph can still be anything.

This is the result of the dataset type being added to the template. Ive already talked to @HowardHuang1 and expressed my preference for a different naming here. We are going to work on improving this in 26.10 for sure.

tfeher · 2026-08-03T11:34:45Z

+      host_index_ = std::make_shared<host_index_type>(cuvs::neighbors::cagra::build(
+        handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(dataset_view_host)));
+      index_ =
+        std::make_shared<index_type>(detail::to_graph_only_index<T, IdxT>(handle_, *host_index_));


What is the differnece betweenhost_index_and index_? Isn't host_index_ already a graph_only index? (due to add_dataset_on_build=false?

host_index_ (host_standard_index) is introduced as a workaround in this PR alongside index_ (device_padded_index, searchable) to support various input dataset types with minimum number of copies.

tfeher · 2026-08-03T11:38:33Z

+    staging = raft::make_device_matrix<T, int64_t>(handle_, rows, dim_);
+    raft::copy(staging.data_handle(),
+               dataset,
+               static_cast<size_t>(rows) * dim_,


If this is called on the whole dataset, then we might run OOM. We should have a vpq quantization method that takes host input data.

HowardHuang1 and others added 2 commits July 31, 2026 04:44

Fix dataset API PR regressions

daaa2f2

achirkin self-assigned this Jul 31, 2026

achirkin added bug Something isn't working non-breaking Introduces a non-breaking change labels Jul 31, 2026

github-project-automation Bot added this to Unstructured Data Processing Jul 31, 2026

Fix copyright

ee5bf13

achirkin added 2 commits July 31, 2026 16:14

Undo unnecessary changes and make diskann copy the dataset from the b…

200c064

…enchmark-provided dataset to avoid ambiguity w.r.t. data residence

Avoid 2x graph on device

59657bb

achirkin changed the title ~~WIP Fix dataset API PR regressions~~ Fix benchmark regressions from #1846 Jul 31, 2026

Avoid a staging full dataset on device when using compressed dataset

42923a4

achirkin marked this pull request as ready for review July 31, 2026 15:41

achirkin requested a review from a team as a code owner July 31, 2026 15:41

cjnolet changed the base branch from release/26.08 to main July 31, 2026 23:01

cjnolet requested review from a team as code owners July 31, 2026 23:01

cjnolet requested a review from msarahan July 31, 2026 23:01

Merge branch 'main' into fix-dataset-api-benchmark-regressions

8be966a

achirkin removed request for a team August 3, 2026 09:59

achirkin removed request for a team and msarahan August 3, 2026 09:59

tfeher requested changes Aug 3, 2026

View reviewed changes

Navigation Menu

Uh oh!

Fix benchmark regressions from #1846 - #2384

Fix benchmark regressions from #1846#2384
achirkin wants to merge 7 commits into
NVIDIA:mainfrom
achirkin:fix-dataset-api-benchmark-regressions

achirkin commented Jul 31, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 31, 2026

Uh oh!

achirkin commented Jul 31, 2026

Uh oh!

achirkin commented Jul 31, 2026

Uh oh!

cjnolet commented Jul 31, 2026 •

edited

Loading

Uh oh!

achirkin commented Jul 31, 2026

Uh oh!

tfeher left a comment

Uh oh!

tfeher Aug 3, 2026

Uh oh!

achirkin Aug 3, 2026

Uh oh!

cjnolet Aug 3, 2026 •

edited

Loading

Uh oh!

tfeher Aug 3, 2026

Uh oh!

achirkin Aug 3, 2026

Uh oh!

tfeher Aug 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

achirkin commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jul 31, 2026

Uh oh!

achirkin commented Jul 31, 2026

Uh oh!

achirkin commented Jul 31, 2026

Uh oh!

cjnolet commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

achirkin commented Jul 31, 2026

Uh oh!

tfeher left a comment

Choose a reason for hiding this comment

Uh oh!

tfeher Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

achirkin Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

cjnolet Aug 3, 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!

tfeher Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

achirkin Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

tfeher Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

achirkin commented Jul 31, 2026 •

edited

Loading

cjnolet commented Jul 31, 2026 •

edited

Loading

cjnolet Aug 3, 2026 •

edited

Loading