Cache IVF-PQ and select-warpsort kernel launch parameters to reduce latency - #1786
Conversation
achirkin
added
improvement
a653b8b to
7aefb3b
Compare
|
@achirkin can you retarget this PR to 24.4? It's too late to get this into 24.02. |
|
I've benchmarked prims::SelectK, prims:KNN/ivf_pq (with n_queries set to 1), and ann-bench ivf-pq. Performance changes in the prims benchmarks are very minimal; most of the cases are the same withing +-2-3% time, a few cases became faster by 5-8%. The most benefit comes in ann-bench run in the "throughput" mode. At 4 threads, it's 20%, at 32 threads it's 60% faster. |
achirkin
added
3 - Ready for Review
and removed
2 - In Progress
tfeher
left a comment
tfeher
left a comment
There was a problem hiding this comment.
Thanks for this PR Artem, it looks good to me. I have one question, but that need not block merging this.
| inline auto operator()(const search_kernel_key& x) const noexcept -> std::size_t | ||
| { | ||
| return (size_t{x.manage_local_topk} << 63) + | ||
| size_t{x.topk} * size_t{x.n_probes} * size_t{x.n_queries} + |
There was a problem hiding this comment.
Why is this hash with parameters ("*") multiplied used? I guess occasional collisions do not matter much for us, but a proper hash_combine functions is not much more complex than this.
There was a problem hiding this comment.
Good point! But does raft depend on boost already? There seem to be no standard library equivalent for this.
Also, in the current implementation of LRU cache, we don't use hash values at all (we changed the implementation from a hashmap to a vector of key-value tuples).
|
/merge |
…atency (NVIDIA#1786) This PR aims at reducing the latency in IVF-PQ and related functions, especially with small work sizes and in the "throughput" benchmark mode. - Add kernel config caching to ivf_pq::search::compute_similarity kernel - Add kernel config caching to select::warpsort - Fix the memory_resource usage in `matrix::select_k`: make sure all temporary allocations use raft's workspace memory resource. Authors: - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Tamas Bela Feher (https://github.com/tfeher) URL: NVIDIA#1786
4 participants
This PR aims at reducing the latency in IVF-PQ and related functions, especially with small work sizes and in the "throughput" benchmark mode.
matrix::select_k: make sure all temporary allocations use raft's workspace memory resource.