[REVIEW] Add device-memory overload for cagra::helpers::optimize - #2423
Open
mnorris11 wants to merge 1 commit into
Open
[REVIEW] Add device-memory overload for cagra::helpers::optimize#2423mnorris11 wants to merge 1 commit into
mnorris11 wants to merge 1 commit into
Conversation
`cagra::helpers::optimize` currently accepts host matrices only, so a caller
that already holds its k-NN graph in device memory -- for example the output of
`all_neighbors::build` -- must copy the graph to host, optimize, and copy back.
The device implementation already exists. `graph::optimize` is templated on both
mdspan accessors and `make_reverse_graph_gpu` has an `is_device_accessible` fast
path (graph_core.cuh:826), and `batch_load_iterator` switches to a zero-copy
`kPassthrough` mode for device accessors. That code is unreachable today because
`detail::optimize` in cagra_build.cuh erases the caller's accessor:
using g_accessor_internal =
raft::host_device_accessor<cuda::std::default_accessor<internal_IdxT>,
raft::memory_type::host>;
so `graph::optimize` is only ever instantiated with host accessors. With host
accessors the reverse-graph phase degrades into `graph_degree` separate host
column gathers, each with its own H2D copy and a full stream synchronisation.
This change:
- propagates the caller's memory types through `detail::optimize` instead of
erasing them, and makes `new_graph` accessor-generic there and in
`cagra::optimize`;
- adds a `device_matrix_view` overload of `cagra::helpers::optimize` to the
public API;
- exposes the existing `guarantee_connectivity` flag on the public overloads,
which previously could not be reached from outside.
Both public overloads keep their existing signatures via a defaulted argument,
so this is source compatible.
Measured on 100M x 129d vectors (graph degree 32, intermediate degree 32) on
8x H100, as part of a multi-GPU CAGRA build: the optimize step goes from 119.1s
to 1.65s, a 72x reduction, with recall unchanged. At that scale it takes the
whole build->serialize pipeline from 8.0 to 5.4 minutes.
Tests: adds a device-to-device case and a case asserting the device overload
produces the same graph as the host overload for the same input.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Disclaimer: AI mostly generated, but validated on a 8x H100 host. Let me know if I should delete the lengthy AI comments or if this PR is silly and there's better workarounds. We will add a hack in Faiss at facebookresearch/faiss#5500 until we can merge this PR / import the latest cuVS version (probably around 26.10? or later?).
made.
Problem
cagra::helpers::optimizecurrently accepts host matrices only, so a caller that already holds its k-NN graph in device memory -- for example the output ofall_neighbors::build-- must copy the graph to host, optimize, and copy back.The device implementation already exists.
graph::optimizeis templated on both mdspan accessors andmake_reverse_graph_gpuhas anis_device_accessiblefast path (graph_core.cuh:826), andbatch_load_iteratorswitches to a zero-copykPassthroughmode for device accessors. That code is unreachable today becausedetail::optimizein cagra_build.cuh erases the caller's accessor:so
graph::optimizeis only ever instantiated with host accessors. With host accessors the reverse-graph phase degrades intograph_degreeseparate host column gathers, each with its own H2D copy and a full stream synchronisation.This change
detail::optimizeinstead of erasing them, and makesnew_graphaccessor-generic there and incagra::optimize;device_matrix_viewoverload ofcagra::helpers::optimizeto the public API;guarantee_connectivityflag on the public overloads, which previously could not be reached from outside.Both public overloads keep their existing signatures via a defaulted argument, so this is source compatible.
Measured on 100M x 128d vectors (graph degree 32, intermediate degree 32) on 8x H100, as part of a multi-GPU CAGRA build: the optimize step goes from 119.1s to 1.65s, a 72x reduction, with recall unchanged.
added.
Tests: adds a device-to-device case and a case asserting the device overload produces the same graph as the host overload for the same input.
noted here: https://help.github.com/articles/closing-issues-using-keywords/
None I think?