Skip to content

perf(rccl): skip the temporary stream when capturing single-stream launches #14118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wenkaidu merged 4 commits into from
Oct 28, 2026
Merged

perf(rccl): skip the temporary stream when capturing single-stream launches #14118

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3f92b5e
perf(rccl): skip the temporary stream when capturing single-stream la…
wenkaidu Oct 21, 2026
7ffef35
test(rccl): document the GRAPH_STREAM_ORDERING=0 capture segfault on …
wenkaidu Oct 21, 2026
a86db9a
fix(rccl): address review comments on the single-stream capture PR
wenkaidu Oct 23, 2026
3cd8361
fix(rccl): address second round of review comments on the single-stre…
wenkaidu Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion projects/rccl/src/enqueue/enqueue.cc
View file
Open in desktop
Comment thread
wenkaidu marked this conversation as resolved.
Resolved
Original file line number Diff line number Diff line change
Expand Up @@ -2896,7 +2896,14 @@ ncclResult_t ncclLaunchFinish(struct ncclComm* comm) {
// So we could do this:
// CUDACHECK(cudaStreamWaitEvent(deviceStream, finishedEvent, 0));
// But instead we do:
NCCLCHECK(ncclStreamAdvanceToEvent(planner->capturingGraph, deviceStream, finishedEvent));
if (capturing && planner->numStreams == 1) {
Comment thread
wenkaidu marked this conversation as resolved.
Resolved
// Nothing is captured on launchStream between recording finishedEvent and here, so
// its capture dependencies are exactly what a wait on finishedEvent would install.
// Reading them directly avoids the temporary stream ncclStreamAdvanceToEvent creates.
Comment thread
pvallem marked this conversation as resolved.
Resolved
NCCLCHECK(ncclStreamAdvanceToStream(deviceStream, launchStream));
} else {
NCCLCHECK(ncclStreamAdvanceToEvent(planner->capturingGraph, deviceStream, finishedEvent));
}
}

// Each userStream[i] waits on userStream[0]
Expand Down
6 changes: 6 additions & 0 deletions projects/rccl/src/include/strongstream.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,12 @@ ncclResult_t ncclStreamWaitStream(cudaStream_t a, cudaStream_t b, cudaEvent_t sc
// Like cudaStreamWaitEvent except `e` must be strictly ahead of everything in `s`.
ncclResult_t ncclStreamAdvanceToEvent(struct ncclCudaGraph g, cudaStream_t s, cudaEvent_t e);

// Capture-only equivalent of ncclStreamAdvanceToEvent for an event just recorded on
// `source`: replaces the capture dependencies of `s` with those of `source`. Both
// streams must be capturing into the same graph, and `source` must be strictly ahead
// of everything in `s`.
ncclResult_t ncclStreamAdvanceToStream(cudaStream_t s, cudaStream_t source);

// Synchrnoization does not need the strong stream to be acquired.
ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss);

Expand Down
64 changes: 36 additions & 28 deletions projects/rccl/src/misc/strongstream.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,50 +367,58 @@ ncclResult_t ncclStreamWaitStream(cudaStream_t a, cudaStream_t b, cudaEvent_t sc
return ncclSuccess;
}

ncclResult_t ncclStreamAdvanceToEvent(struct ncclCudaGraph g, cudaStream_t s, cudaEvent_t e) {
if (g.graphId == ULLONG_MAX) {
CUDACHECK(cudaStreamWaitEvent(s, e, 0));
} else {
cudaStream_t tmp;
CUDACHECK(cudaStreamCreateWithFlags(&tmp, cudaStreamNonBlocking));
CUDACHECK(cudaStreamWaitEvent(tmp, e, 0));

cudaStreamCaptureStatus status;
cudaGraphNode_t const* nodes;
size_t count = 0;
// Replace the capture dependencies of `s` with the current capture dependencies of `source`.
static ncclResult_t setCaptureDependenciesFrom(cudaStream_t s, cudaStream_t source) {
cudaStreamCaptureStatus status;
cudaGraphNode_t const* nodes;
Comment thread
wenkaidu marked this conversation as resolved.
Resolved
size_t count = 0;
#if CUDART_VERSION >= 13000
cudaError_t res = hipStreamGetCaptureInfo_v3(tmp, &status, nullptr, nullptr, &nodes, nullptr, &count);
cudaError_t res = hipStreamGetCaptureInfo_v3(source, &status, nullptr, nullptr, &nodes, nullptr, &count);
#else
cudaError_t res = hipStreamGetCaptureInfo_v2(tmp, &status, nullptr, nullptr, &nodes, &count);
cudaError_t res = hipStreamGetCaptureInfo_v2(source, &status, nullptr, nullptr, &nodes, &count);
#endif

#if CUDART_VERSION >= 12030
if (res == cudaErrorLossyQuery) { // CUDA is telling us the dependencies have edge annotations.
cudaGraphEdgeData const* edges;
CUDACHECK(cudaStreamGetCaptureInfo_v3(tmp, &status, nullptr, nullptr, &nodes, &edges, &count));
CUDACHECK(cudaStreamUpdateCaptureDependencies_v2(s, (cudaGraphNode_t*)nodes, edges, count,
cudaStreamSetCaptureDependencies));
}
if (res == cudaErrorLossyQuery) { // CUDA is telling us the dependencies have edge annotations.
cudaGraphEdgeData const* edges;
CUDACHECK(cudaStreamGetCaptureInfo_v3(source, &status, nullptr, nullptr, &nodes, &edges, &count));
CUDACHECK(cudaStreamUpdateCaptureDependencies_v2(s, (cudaGraphNode_t*)nodes, edges, count,
cudaStreamSetCaptureDependencies));
}
#else
if (false) {
}
if (false) {
}
#endif
else {
CUDACHECK(res /* = cudaStreamGetCaptureInfo_v2(...)*/);
else {
CUDACHECK(res /* = cudaStreamGetCaptureInfo_v2(...)*/);
#if CUDART_VERSION >= 13000
CUDACHECK(cudaStreamUpdateCaptureDependencies_v2(s, (cudaGraphNode_t*)nodes, nullptr, count,
cudaStreamSetCaptureDependencies));
CUDACHECK(cudaStreamUpdateCaptureDependencies_v2(s, (cudaGraphNode_t*)nodes, nullptr, count,
cudaStreamSetCaptureDependencies));
#else
CUDACHECK(cudaStreamUpdateCaptureDependencies(s, (cudaGraphNode_t*)nodes, count,
cudaStreamSetCaptureDependencies));
CUDACHECK(cudaStreamUpdateCaptureDependencies(s, (cudaGraphNode_t*)nodes, count,
cudaStreamSetCaptureDependencies));
#endif
}
}
return ncclSuccess;
}

ncclResult_t ncclStreamAdvanceToEvent(struct ncclCudaGraph g, cudaStream_t s, cudaEvent_t e) {
if (g.graphId == ULLONG_MAX) {
CUDACHECK(cudaStreamWaitEvent(s, e, 0));
} else {
cudaStream_t tmp;
CUDACHECK(cudaStreamCreateWithFlags(&tmp, cudaStreamNonBlocking));
CUDACHECK(cudaStreamWaitEvent(tmp, e, 0));
NCCLCHECK(setCaptureDependenciesFrom(s, tmp));
CUDACHECK(cudaStreamDestroy(tmp));
}
return ncclSuccess;
}

ncclResult_t ncclStreamAdvanceToStream(cudaStream_t s, cudaStream_t source) {
Comment thread
wenkaidu marked this conversation as resolved.
Resolved
return setCaptureDependenciesFrom(s, source);
}

ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss) {
#if ROCM_VERSION >= 60100
CUDACHECK(cudaStreamWaitEvent(ss->liveStream, ss->serialEvent, 0));
Expand Down
1 change: 1 addition & 0 deletions projects/rccl/test/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -287,6 +287,7 @@ if(BUILD_TESTS)
DoneEventOrderingTests.cpp
Fp8VerifiableTests.cpp
GatherTests.cpp
GraphCaptureStreamOrderingTests.cpp
Comment thread
wenkaidu marked this conversation as resolved.
Resolved
GroupCallTests.cpp
NonBlockingTests.cpp
RasJsonTests.cpp
Expand Down
Loading
Loading