Fix stream ordering for brute-force L2 sqrt post-processing

[dantegd]

Run the fused brute-force KNN distance post-processing on the same selected CUDA stream as fusedL2Knn.

Problem

In brute_force_knn_impl(), a stream is selected using:

auto stream = raft::resource::get_next_usable_stream(handle, i);

fusedL2Knn is launched on this selected stream. For L2SqrtExpanded and L2SqrtUnexpanded, the subsequent square-root map is instead launched using the original handle, whose main stream may differ from the selected stream:

fusedL2Knn(..., stream, ...);

raft::linalg::map(
  handle,
  ...);

There is no dependency making the main stream wait for the selected stream before the map reads the KNN output. The later stream-pool synchronization does not repair this read-before-write race.

Under concurrent cuML HDBSCAN workloads using managed memory on GB300, the corrupted CUDA context was subsequently reported by RAFT MST as:

copy_if failed on 2nd step: cudaErrorIllegalAddress

Related cuML issue: NVIDIA/cuml#8510
Companion RAFT PR: NVIDIA/raft#3125

[betatim]

Do we need more review on this or can we merge it?

Getting this merged would (hopefully) make the cuml nightly CI happy

[achirkin]

/merge

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/cuvs/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 1a14e2fd-4c2c-4adc-881f-48efc9018ee1

📥 Commits

Reviewing files that changed from the base of the PR and between c525793 and 9bf85cf.

📒 Files selected for processing: 1

• cpp/src/neighbors/detail/knn_brute_force.cuh

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Summary

Summary by CodeRabbit

• Bug Fixes
  • Improved fused L2 nearest-neighbor processing by ensuring its post-processing uses the active execution stream, helping keep results consistent with the selected stream.

The fused L2 kNN path now uses a copied raft::resources handle configured with the current stream-pool stream for result post-processing.

Changes: L2 kNN stream selection

Layer / File(s)	Summary
Configure stream for post-processing cpp/src/neighbors/detail/knn_brute_force.cuh	The L2 result post-processing map uses a copied resources handle set to the current stream-pool stream.

Priority: ⬆️ High

Estimated code review effort: 2 (Simple) | ~5 minutes

Change: Bug fix

Suggested reviewers: achirkin, divyegala, lowener

Merge Risk: ⚪ Minimal · up to 9bf85

L2 square-root post-processing now runs on the same CUDA stream as the fused kNN work, preserving stream ordering for supported calls. No actionable merge risk remains.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: correcting stream ordering for brute-force L2 square-root post-processing.
Description check		The description is directly related to the changeset. It explains the stream mismatch, the resulting race condition, and the intended fix.
Docstring Coverage		No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0 files. (1 skipped: 1 unsupported.)

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}