Iterative CAGRA-Q - #1810
Conversation
|
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. |
aamijar
added
non-breaking
… 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.
8fc10ac to
4f4068a
Compare
…a/cuvs into iterative_cagra_q
|
/ok to test 130f586 |
|
/ok to test |
@tarang-jain, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
|
/ok to test 974f920 |
|
/ok to test 248e028 |
| * `n_rows` is `IdxT` and the remaining five are `uint32_t`. | ||
| */ | ||
| template <typename DataT, typename IdxT> | ||
| void serialize_vpq(raft::resources const& res, |
There was a problem hiding this comment.
It seems as though we are supporting PQ dataset serialization in this PR?
So we are supporting serialization and deserialization of index + attached PQ dataset in this PR?
| using index_dataset_view_t = std::remove_cvref_t<decltype(idx.dataset())>; | ||
| if constexpr (cuvs::neighbors::is_vpq_dataset_view_v<index_dataset_view_t>) { | ||
| RAFT_FAIL( | ||
| "CAGRA index serialization is not supported for VPQ indices"); |
There was a problem hiding this comment.
I think we need to fix this for the C API? We are supporting index serialization in CPP but not in C?
Not sure if C API serialization is necessary for this PR.
| } | ||
|
|
||
| template <typename T> | ||
| static auto make_vpq_from_dense_dataset(raft::resources* res_ptr, |
There was a problem hiding this comment.
Making a note to rename any new instances of VPQ to PQ instead.
The existing instances will be renamed as part of #2494.
| -> index<uint8_t, uint32_t, device_vpq_dataset_view<half, int64_t>>; | ||
|
|
||
| auto update_dataset(raft::resources const& res, | ||
| vpq_f16_index<float>&& cagra_index, |
There was a problem hiding this comment.
We have a bunch of uses of the alias vpq_f16_index. I'm wondering whether we should change this alias to just be pq_index as part of this PR.
Or use the full namespace version:
cuvs::neighbors::cagra::index<
float, // query/input type
uint32_t, // neighbor ID type
cuvs::neighbors::device_vpq_dataset_view<half, int64_t> // dataset view
>
It is a bit verbose though.
| auto build(raft::resources const& res, | ||
| const cuvs::neighbors::cagra::index_params& params, | ||
| cuvs::neighbors::device_vpq_dataset_view<half, int64_t> const& dataset) | ||
| -> index<float, uint32_t, cuvs::neighbors::device_vpq_dataset_view<half, int64_t>>; |
There was a problem hiding this comment.
I just noticed we are supporting half but not float for vpq datasets. I think we do have a float version of vpq? Is that something we are supporting in Iterative CAGRA-Q? I think CAGRA-Q search kernels already support it?
| device_vpq_dataset_view<half, int64_t> dataset) -> vpq_f16_index<int8_t>; | ||
| auto update_dataset(raft::resources const& res, | ||
| vpq_f16_index<uint8_t>&& cagra_index, | ||
| device_vpq_dataset_view<half, int64_t> dataset) -> vpq_f16_index<uint8_t>; |
There was a problem hiding this comment.
We definitely need to add the overloads to go from PQ backed index to Dense backed index.
That is the main use case for us too. As seen from benchmarking searching on dense backed index is what gives us 99%+ recall potential.
| size_t size_; | ||
| }; | ||
| template <typename T, typename MathT, typename IdxT> | ||
| void reconstruct_vpq_queries(raft::resources const& res, |
There was a problem hiding this comment.
Thinking longer term, have we considered using the reconstruction functionality from the existing PQ APIs? I don't think it can be directly replaced right now because of differing format, but perhaps something to explore later.
8 participants
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.