Skip to content
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

Merged
merged 8 commits into from
Jan 23, 2024

Conversation

benfred
Copy link
Member

@benfred benfred commented Jan 16, 2024

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.

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.
@benfred benfred requested a review from a team as a code owner January 16, 2024 18:44
@github-actions github-actions bot added the cpp label Jan 16, 2024
@benfred benfred changed the title 0Allow topk larger than 1024 in CAGRA Allow topk larger than 1024 in CAGRA Jan 16, 2024
@benfred benfred added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Jan 16, 2024
Copy link
Contributor

@tfeher tfeher left a 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.

Copy link
Contributor

@tfeher tfeher left a 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!

@cjnolet
Copy link
Member

cjnolet commented Jan 23, 2024

/merge

@rapids-bot rapids-bot bot merged commit 0586fc3 into rapidsai:branch-24.02 Jan 23, 2024
61 checks passed
abc99lr pushed a commit to abc99lr/raft that referenced this pull request Jan 23, 2024
abc99lr pushed a commit to abc99lr/raft that referenced this pull request Jan 23, 2024
cjnolet added a commit to cjnolet/raft that referenced this pull request Feb 8, 2024
raydouglass pushed a commit that referenced this pull request Feb 9, 2024
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)
benfred added a commit to benfred/raft that referenced this pull request Feb 13, 2024
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.
rapids-bot bot pushed a commit that referenced this pull request Feb 14, 2024
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
cpp improvement Improvement / enhancement to an existing function non-breaking Non-breaking change
Projects
Development

Successfully merging this pull request may close these issues.

3 participants