ScaNN: Fix AVQ prefetch

[rmaschal]

Switching to usage of modern RAFT in cuVS (#1837) introduced a bug where the prefetched gather for AVQ is performed using the stream associated with raft::device_resources rather than the provided stream for copying. This led to two issues:

1. Elimination of the benefit for prefetching, as copies where scheduled on the same stream as other gpu work
2. Possible recall loss. Synchronization was still performed against the copy stream, potentially allowing host to proceed before the prefetch copy is complete.

This PR sets the stream associated with the resource to the copy stream before prefetching, and back when done.

[aamijar]

Good catch @rmaschal, and LGTM!

[divyegala]

You can create a new resources object for no cost:

Suggested change

	// For prefetching to overlap with other gpu work
	// we need to schedule copies on the provided copy stream stream_
	auto copy_res = raft::resources(stream_);
	// htod
	auto h_cluster_ids =
	raft::make_pinned_vector_view<LabelT, int64_t>(cluster_ids_buf_.data_handle(), size);
	@@ -532,6 +537,8 @@ class cluster_loader {
	raft::copy(copy_res, cluster_vectors, raft::make_const_mdspan(pinned_cluster));
	raft::resource::sync_stream(res, stream_);

[achirkin]

I vote on this too. Although not very relevant here, another benefit of creating a temporary resources handle is the automatic reset of the changes if an exception is raised. However, I'd like to warn that there are caveats in doing this:

1. I'd suggest to use a copy the resources handle rather than creating a new resources handle to inherit all other initialized resources from the current handle.

auto copy_res = raft::resources(res);
raft::resource::set_cuda_stream(res, stream);

2. Even when you do a copy: if you've used any other not-yet-initialized resources from the copy_res handle, it would create those resources on every invocation of prefetch_cluster, which would incur significant overhead.

[rmaschal]

Thank you for the comments and info. I now copy the passed resources, set the stream, and use that for scheduling copies

[aamijar]

Hi @rmaschal, is this targeting release/26.04? If so, please retarget from main to release/26.04

[rmaschal]

Hi @rmaschal, is this targeting release/26.04? If so, please retarget from main to release/26.04

Hi @aamijar, just saw the comment, I assume it's too late for 26.04? In that case main is fine, then targeting 26.06

[aamijar]

Hi @rmaschal, looks like we are good to merge this right, or are you seeing CI failures related this PR? No need to keep merging in main, as long as the Recently Updated check is passing. You can also rerun failed jobs if they are flaky.

[rmaschal]

@aamijar yes should be good, some python CI failures but they should be unrelated.

[aamijar]

Great! I will try to rebase to release/26.04. Would be nice to have in the release since it a bug fix/ perf regression fix.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[aamijar]

/ok to test d5d3316

[aamijar]

/merge