[REVIEW] Add device-memory overload for cagra::helpers::optimize - #2423
[REVIEW] Add device-memory overload for cagra::helpers::optimize#2423mnorris11 wants to merge 2 commits into
Conversation
mfoerste4
left a comment
mfoerste4
left a comment
There was a problem hiding this comment.
@mnorris11 , thanks for the PR. Besides the lengthy AI comments you already mentioned the PR looks good to me. I would like to wait for other peoples comments on exposing the mst graph option in its current state though. CC @cjnolet
| raft::host_matrix_view<uint32_t, int64_t, raft::row_major> knn_graph, | ||
| raft::host_matrix_view<uint32_t, int64_t, raft::row_major> new_graph); | ||
| raft::host_matrix_view<uint32_t, int64_t, raft::row_major> new_graph, | ||
| bool guarantee_connectivity = false); |
There was a problem hiding this comment.
Is this parameter really required for your workflow? The mst path should be functionally correct, but is not optimized wrt. speed and memory consumption. I would rather not see it in the public API. @achirkin , what is your opinion on this?
There was a problem hiding this comment.
I don't think this parameter will be a big problem since the default is false, but I agree unless you sure you need this feature it will be easier for us to support the API without it.
There was a problem hiding this comment.
I do absolutely want the guarantee_connectivity to be exposed to the public. The reason this parameter was introduced was not for testing, but rather because we came to realization that some types of neighborhood graphs, such as the RNG graph which underlines assumptions in Vamana, implicitly assume the graph is connected. This can obviously greatly impact the navigability since NSG-style graphs tend to start with an all-neighbors knn graph (which is not guaranteed to be connected).
`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.
1ca817d to
1fe0cff
Compare
cjnolet
added
improvement
4 participants
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?