Fix missing sync_stream in ScaNN build #1224
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
github-project-automation
bot
moved this to Todo
in Vector Search, ML, & Data Mining Release Board
cjnolet
added
improvement
cjnolet
moved this from Todo
to In Progress
in Vector Search, ML, & Data Mining Release Board
cjnolet
approved these changes
Aug 6, 2025
tarang-jain
reviewed
Aug 6, 2025
| dataset_vec_batches.prefetch_next_batch(); | ||
|
|
||
| // Make sure work on device is finished before swapping buffers | ||
| raft::resource::sync_stream(res); |
Contributor
There was a problem hiding this comment.
prefetch_next_batch seems to be synchronizing at the very end. Would it make better sense to have the sync_stream before prefetching the batch (above line 175)?
Contributor
Author
There was a problem hiding this comment.
prefetch_next_batch() only syncs on the stream used for copying, so if the copy happens on a separate stream from the work (e.g. if you want to prfetch/overlap) then work might not be done when buffers are swapped (which is what I am trying to fix). Moving sync stream before prefetch will fix the problem, but at the cost of no overlapping of work/copy, reducing performance.
Contributor
There was a problem hiding this comment.
Good point. I had not realized that prefetching is done using a different stream
tarang-jain
approved these changes
Aug 6, 2025
Member
|
/merge |
github-project-automation
bot
moved this from In Progress
to Done
in Vector Search, ML, & Data Mining Release Board
lowener
pushed a commit
to lowener/cuvs
that referenced
this pull request
Aug 11, 2025
As title. For larger number of clusters (or faster PCI-E), the work required to find cluster assignments would take longer than the HtoD copy for each batch. This would lead to overwriting of one of the batch_load_iterator device buffers before kmeans predict completed, and thus lower recall. This adds a raft::resource::sync_stream at the end of the kmeans predict loop to avoid this. Although not technically necessary (as the prior DtoH copy into pageable memory is synchronous), sync_stream is also added to the end of the quantization loop for consistency/correctness. Also removed one unused comment. Authors: - https://github.com/rmaschal Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Tarang Jain (https://github.com/tarang-jain) URL: rapidsai#1224
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
As title. For larger number of clusters (or faster PCI-E), the work required to find cluster assignments would take longer than the HtoD copy for each batch. This would lead to overwriting of one of the batch_load_iterator device buffers before kmeans predict completed, and thus lower recall.
This adds a raft::resource::sync_stream at the end of the kmeans predict loop to avoid this. Although not technically necessary (as the prior DtoH copy into pageable memory is synchronous), sync_stream is also added to the end of the quantization loop for consistency/correctness.
Also removed one unused comment.