Skip to content

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

Merged
rapids-bot[bot] merged 19 commits into
NVIDIA:release/26.10from
tarang-jain:codex/cuvs-bench-prefetch-async-rmm
Sep 11, 2026
Merged

[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers#2526
rapids-bot[bot] merged 19 commits into
NVIDIA:release/26.10from
tarang-jain:codex/cuvs-bench-prefetch-async-rmm

Conversation

@tarang-jain

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

Copy link
Copy Markdown
Contributor
  • remove stream pool creation from individual wrappers

@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).

Comment thread cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h Outdated
Comment thread cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h Outdated
Comment thread cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h Outdated
@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
Comment thread cpp/bench/ann/src/common/util.hpp
@tarang-jain
tarang-jain requested a review from a team as a code owner September 3, 2026 20:49

@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 for the updates, LGTM!
But why did you undo the async managed resource change? I would expect it to slightly improve performance during build.

@tarang-jain

Copy link
Copy Markdown
Contributor Author

But why did you undo the async managed resource change? I would expect it to slightly improve performance during build.

It was giving problems (test failures) with CUDA 12.

@achirkin

achirkin commented Sep 5, 2026

Copy link
Copy Markdown
Contributor

It was giving problems (test failures) with CUDA 12.

Omg, how would it cause test failures if it's in the benchmark component?

msarahan and others added 3 commits September 8, 2026 12:07
We already have computed the centroid assignment from the previous set of candidates in the previous iteration. We do not need the full PW reduction in the current iteration, but only across the newly added candidates.

Profile on A6000 with changes from NVIDIA#2249:
Case 1: 1M * 1536 dataset, with 1024 centroids. Here we profile only the scalable kmeans++ initialization
1. NVIDIA#2249 alone:

Time | Total Time | Instances | Avg | Med | Min | Max | StdDev | Name
-- | -- | -- | -- | -- | -- | -- | -- | --
86.2% | 18.815 s | 16 | 1.176 s | 265.906 ms | 442.142 μs | 4.202 s | 1.493 s | fused_1nn_f_i32_relaxed
11.2% | 2.441 s | 1023 | 2.386 ms | 2.499 ms | 2.163 ms | 2.615 ms | 158.442 μs | cutlass gemm

2. NVIDIA#2249 + this PR:

Time | Total Time | Instances | Avg | Med | Min | Max | StdDev | Name
-- | -- | -- | -- | -- | -- | -- | -- | --
72.0% | 7.994 s | 15 | 532.946 ms | 511.092 ms | 434.372 μs | 4.213 s | 1.052 s | fused_1nn_f_i32_relaxed
22.9% | 2.538 s | 1023 | 2.481 ms | 2.497 ms | 2.118 ms | 3.124 ms | 179.081 μs | cutlass gemm

Authors:
  - Tarang Jain (https://github.com/tarang-jain)
  - Divye Gala (https://github.com/divyegala)

Approvers:
  - Victor Lafargue (https://github.com/viclafargue)

URL: NVIDIA#2438
@tarang-jain
tarang-jain changed the base branch from main to release/26.08 September 8, 2026 20:24
@tarang-jain
tarang-jain requested review from a team as code owners September 8, 2026 20:24
@tarang-jain
tarang-jain requested a review from msarahan September 8, 2026 20:24
@tarang-jain

Copy link
Copy Markdown
Contributor Author

Omg, how would it cause test failures if it's in the benchmark component?

Basically the async managed memory resource is only available in CUDA 13+.
https://github.com/rapidsai/rmm/blob/1de695335e8585d17f916e612e6f69a623ca1a9d/cpp/src/mr/detail/cuda_async_managed_memory_resource_impl.cpp#L23
The tests that I was talking about are the python cuvs-bench tests.

@tarang-jain
tarang-jain changed the base branch from release/26.08 to release/26.10 September 8, 2026 20:27
@tarang-jain
tarang-jain removed request for a team and msarahan September 8, 2026 20:28
@tarang-jain
tarang-jain removed request for a team September 9, 2026 18:10
@tarang-jain

Copy link
Copy Markdown
Contributor Author

/merge

@rapids-bot
rapids-bot Bot merged commit c4e23e4 into NVIDIA:release/26.10 Sep 11, 2026
184 of 187 checks passed
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.

4 participants