CUDA Graph Capture

cuDecomp transpose and halo operations can be captured as part of a caller-owned CUDA Graph. This allows an application to capture a larger workflow containing cuDecomp operations and replay it with the standard CUDA Graph APIs.

Compatibility

External capture compatibility

Communication family	Supported	Notes
MPI	No	Performs host-side MPI communication and CUDA stream synchronization.
NCCL	Yes	Includes the pipelined transpose backend.
NVSHMEM	Yes	Includes the pipelined and device-initiated transpose backends; requires NVSHMEM 2.9 or newer. NVSHMEM 2.6 through 2.8 remain supported for normal, uncaptured operation.

External capture is separate from cuDecomp’s internal graph optimization. The CUDECOMP_ENABLE_CUDA_GRAPHS environment variable must be disabled. If it is enabled, a cuDecomp operation called during external capture returns CUDECOMP_RESULT_NOT_SUPPORTED. Performance reporting must also be disabled during external capture.

Workspace preparation

Caller-owned workspace can be used directly during capture. The workspace must remain valid until all graph executables and graphs that reference it have been destroyed.

Automatic workspace requires a completed warmup before capture. Run each operation once, using the same datatype and configuration that will be captured, and synchronize its stream before beginning capture. The warmup allocates enough workspace and establishes any required communication-buffer registrations. An operation captured before its automatic workspace is ready returns CUDECOMP_RESULT_NOT_SUPPORTED.

The automatic workspace allocation domain is frozen when it is first used during capture. A later operation that would require the allocation to grow returns CUDECOMP_RESULT_NOT_SUPPORTED. Warm up the largest operation that the handle will execute before capturing any operation, or use a separate handle for operations with different workspace requirements.

Execution and lifetime requirements

All participating ranks must capture the same collective operation and must consistently use caller-owned or automatic workspace. Operations using the same cuDecomp handle must not overlap in execution, including eager operations and graph replays. Order them on the same stream or with explicit CUDA dependencies. Applications that require concurrent cuDecomp operations must use separate handles.

Complete all graph launches and destroy graph executables and graphs before destroying the associated grid descriptor or finalizing its cuDecomp handle.

Example

The following example warms up an automatically managed workspace, captures one transpose, and instantiates the resulting graph. Error checking is omitted for brevity.

cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO,
                      CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, stream);
cudaStreamSynchronize(stream);

cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal);
cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO,
                      CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, stream);
cudaStreamEndCapture(stream, &graph);

cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0);
cudaGraphLaunch(graph_exec, stream);
cudaStreamSynchronize(stream);

cudaGraphExecDestroy(graph_exec);
cudaGraphDestroy(graph);