Cache IVF-PQ and select-warpsort kernel launch parameters to reduce latency #1786
Conversation
achirkin
added
improvement
achirkin
force-pushed
the
fea-cache-ivf-pq-params
branch
from
a653b8b to
7aefb3b
Compare
achirkin
added
2 - In Progress
|
@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
| 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 |
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.