-
Notifications
You must be signed in to change notification settings - Fork 197
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Allow topk larger than 1024 in CAGRA #2097
Conversation
This change allows CAGRA search to have an arbitrarily large top-k, instead of being limited to 1024 like in the previous code. This works by using the multi-kernel search path, and replacing the _cuann_find_topk code with the matrix::select_k code - which can handle large K values.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @benfred for the PR! Overall it looks good: while the extra copies are not ideal (and we should add it to a tracker to improve that), it alleviates the max-k problem without affecting the perf where k< 1024. Just one comment below.
cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @benfred for the update, LGTM!
/merge |
This reverts commit 0586fc3.
This reverts commit fdd4ad2.
This reverts commit 0586fc3.
RAFT C++ tests were not running for a portion of the 24.02 development cycle, until the merger of rapidsai/rapids-cmake#533. This PR fixes some failing tests and reverts PRs that caused test failures that were silent until now, specifically #2097 and #2085. These features will be revisited in a subsequent release. Authors: - Malte Förster (https://github.com/mfoerste4) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Ben Frederickson (https://github.com/benfred) - Bradley Dice (https://github.com/bdice)
Reapply changes from rapidsai#2097 This change allows CAGRA search to have an arbitrarily large top-k, instead of being limited to 1024 like in the previous code. This works by using the multi-kernel search path, and replacing the _cuann_find_topk code with the matrix::select_k code - which can handle large K values.
Reapply changes from #2097 This change allows CAGRA search to have an arbitrarily large top-k, instead of being limited to 1024 like in the previous code. This works by using the multi-kernel search path, and replacing the _cuann_find_topk code with the matrix::select_k code - which can handle large K values. Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #2181
This change allows CAGRA search to have an arbitrarily large top-k, instead of being limited to 1024 like in the previous code.
This works by using the multi-kernel search path, and replacing the _cuann_find_topk code with the matrix::select_k code - which can handle large K values.