[FEA] IVF-Flat optimize loading cluster data for large batch search
Nobody has claimed this yet.
- Dominant language
- Cuda
- Stars
- 1k
- Forks
- 251
- Avg merge
- 1d 8h
- Merged PRs (30d)
- 31
Description
Is your feature request related to a problem? Please describe.
During IVF-Flat search a query vector is compared to all the vectors from n_probes clusters, and we have n_queries * n_probes query-probe pairs. For large batch search, when n_queries * n_probes > n_clusters then there will be clusters that are compared to more than one query vector.
The execution time of IVF-Flat is determined by the time to load the clusters from memory. Currently the query-probe pairs are sorted according to query index. To improve memory load time, we can sort the query-probe pairs according to the probe id (cluster label).
Describe the solution you'd like
Sort the query-probe pairs during fine search for better cache reuse. This is already implemented for IVF-PQ, and the same can be applied for IVF-Flat as well:
https://github.com/rapidsai/raft/blob/734298013b02b43d57275b342ab39d1dfd102543/cpp/include/raft/neighbors/detail/ivf_pq_search.cuh#L529-L569
Additional context
In IVF-Flat search we typically have 0.1-1% of the clusters searched, therefore this optimization is expected to help with batch size that is correspondingly large (hundreds or thousends of query vectors). We have a helper utility to calculate the expected number of times a cluster is loaded. This can be used to decide whether to sort the input data or not.
https://github.com/rapidsai/raft/blob/734298013b02b43d57275b342ab39d1dfd102543/cpp/include/raft/neighbors/detail/ivf_pq_search.cuh#L396
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in cpp/include/raft/neighbors/detail/ivf_flat_search-inl.cuh around the query-probe sorting near line 175. Compare the IVF-PQ implementation at cpp/include/raft/neighbors/detail/ivf_pq_search.cuh lines 529-569 and review the cluster-load helper near line 396. Done means IVF-Flat large-batch searches can reuse cluster loads and show the expected performance improvement.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 38/100