Skip to content

[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers - #2526

Open
tarang-jain wants to merge 6 commits into
NVIDIA:mainfrom
tarang-jain:codex/cuvs-bench-prefetch-async-rmm
Open

[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers#2526
tarang-jain wants to merge 6 commits into
NVIDIA:mainfrom
tarang-jain:codex/cuvs-bench-prefetch-async-rmm

Conversation

@tarang-jain

@tarang-jain tarang-jain commented Aug 29, 2026

Copy link
Copy Markdown
Contributor
  • Cleanup an old non_blocking_stream struct
  • remove stream pool creation from individual wrappers
  • Make the large memory resource allocations async

@tarang-jain
tarang-jain requested a review from a team as a code owner August 29, 2026 00:01
@tarang-jain tarang-jain added non-breaking Introduces a non-breaking change improvement Improves an existing functionality labels Aug 29, 2026
@achirkin

Copy link
Copy Markdown
Contributor

Could you please give a bit more detail on where the blocking happens? We've been trying to avoid using default streams across raft and cuvs for a long time. If they still appear somewhere within the library, it probably makes sense to fix there?

@achirkin achirkin left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I think the change makes perfect sense and can only improve the performance. We don't need any extra stream syncs, because the replaced managed memory was used as a large workspace resource (which is async by contract).
Just please update the description with the explanation of the problem you wanted to solve and do a small cleanup (see below).

size_t n_streams = cagra_indices.size();
raft::resource::set_cuda_stream_pool(composite_handle,
std::make_shared<rmm::cuda_stream_pool>(n_streams));
make_non_blocking_stream_pool(n_streams));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit surprised how this stream pool setting has appeared in all wrapper code in the first place.
Please remove all these repeated set_cuda_stream_pools from the wrapper files and put it exactly once either in the shared_raft_resources or in configured_raft_resources object (where more appropriate).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I got rid of this behavior in the other ivf wrappers. But for cagra it seems this particular pool is a special case, for the composite index.

Comment thread cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h Outdated
configured_raft_resources() : configured_raft_resources{std::make_shared<shared_raft_resources>()}
{
raft::resource::set_cuda_stream_pool(
*res_, std::make_shared<rmm::cuda_stream_pool>(1, rmm::cuda_stream::flags::non_blocking));

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a prefetch stream in the default constructor.

@tarang-jain tarang-jain changed the title [ANN_BENCH] Use Async MR and Non Blocking Streams [ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants