Synchronize MST edge count update

[bdice]

Synchronize the MST stream after asynchronously copying the current edge count from a loop-local host value into prev_mst_edge_count.

rmm::device_scalar::set_value_async() requires its host source to remain alive until the stream is synchronized. Without the synchronization, the copied edge count can be corrupted after the loop iteration ends and later used as the output offset for thrust::copy_if.

This addresses the illegal-address failure observed in the cuGraph MST test on GB300: https://github.com/rapidsai/cugraph/actions/runs/33161386555/job/98885855224#step:12:3056

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[bdice]

AI Agent Analysis

I verified that #3125 is already present in the failing environment and fixes a different stream-ordering bug:

• The cuGraph job linked libraft-26.10.00a32-...-3a306e53.
• 3a306e53 has 06cb6935 (the merge commit for Fix curand stream ordering in MST alteration #3125) as its parent.
• The installed source therefore already contains curandSetStream(randGen, stream).

The remaining lifetime violation is in MST_solver::solve():

auto curr_mst_edge_count = mst_edge_count.value(stream);
...
prev_mst_edge_count.set_value_async(curr_mst_edge_count, stream);

curr_mst_edge_count is local to the for body. rmm::device_scalar::set_value_async() starts an asynchronous H2D copy and requires the referenced host value to remain alive until the stream is synchronized. The loop iteration ends immediately after this call, so the source may be destroyed or its stack storage reused before the copy completes.

On the following iteration, append_src_dst_pair() reads prev_mst_edge_count and constructs its output iterator from it:

mst_src + curr_mst_edge_count
mst_dst + curr_mst_edge_count
mst_weights + curr_mst_edge_count

A corrupted previous edge count therefore becomes an out-of-bounds output offset for thrust::copy_if, matching the cuGraph failure:

copy_if failed on 2nd step: cudaErrorIllegalAddress

The host COO vectors in the cuGraph test remain in scope, and that test performs a device synchronization before calling MST, so the input transfer is not the lifetime problem. The other set_value_async() in RAFT MST uses true_val, which remains alive throughout label_prop() and is followed by synchronizing done.value(stream) calls before the function returns.

This change synchronizes immediately after updating prev_mst_edge_count, satisfying RMM's host-lifetime contract before the loop-local value leaves scope. The PR is left as a draft pending GB300 validation.

Human evaluation from @bdice: Yes, this seems like a real bug and this seems like a proper fix.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7f015b13-fb85-43a3-9b60-bf61113ec9b1

📥 Commits

Reviewing files that changed from the base of the PR and between c7e7678 and 368e42d.

📒 Files selected for processing (1)

• cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved synchronization during minimum spanning tree solver iterations for more reliable results and execution.

Walkthrough

The MST solver adds the CUDA stream resource header and synchronizes the configured resource handle and stream after updating the previous edge count in each solve iteration.

Changes

MST solver synchronization

Layer / File(s)	Summary
Synchronize MST solve iterations cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh	The solve loop calls resource::sync_stream(handle, stream) after storing the current edge count. The file includes the required CUDA stream resource header.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Merge Risk: ⚪ Minimal · up to 368e4

This localized change synchronizes the MST stream before the asynchronously copied edge count is reused; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: achirkin, dantegd, divyegala

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly explains the stream synchronization change, the host-lifetime issue with rmm::device_scalar::set_value_async(), and the resulting illegal-address failure.
Title check	✅ Passed	The title clearly and concisely identifies the main change: synchronizing the MST edge count update.
Docstring Coverage	✅ Passed	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	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	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 PR with unit tests

---

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

[divyegala]

/merge