[maleadt]

CUDA graphs let you record a sequence of GPU operations once and then launch the whole sequence with a single, cheap API call. For code made of many short kernels, which is common in Julia because every broadcast is a kernel, the CPU cost of launching each operation often exceeds the time the GPU spends executing it, so graphs can make such code several times faster. CUDA.jl has had capture, instantiate and @captured for a while, but they didn't work with CUDA.jl's GC-based memory management:

• A graph didn't keep the arrays it used alive. Replaying it after the GC freed one of them was a use-after-free.
• Arrays allocated during capture became graph-owned allocations. The GC then freed them with an error (Graph memory management #1230), and the graph couldn't be launched twice.
• Scalars passed to libraries like cuBLAS were captured as copies from a temporary host buffer, so replays read garbage (@captured with scalar arguments fails with CUDA 5.7 #2691).
• capture disabled the GC for the whole process, and freeing memory from another thread could still invalidate the capture (Freeing memory from finalizers during graph capture #3310).

This PR reworks the integration so that graphs are safe to use with ordinary Julia code, and adds an API for both simple capture-and-replay and for constructing graphs by hand. It builds on the deferred release of resources from #3318, which is what makes capture compatible with the GC. The commits are self-contained steps, and are best reviewed one at a time.

using CUDA, LinearAlgebra

# a small MLP, evaluated in place
function forward!(hs, Ws, bs)
    for i in eachindex(Ws)
        mul!(hs[i+1], Ws[i], hs[i])
        hs[i+1] .= max.(hs[i+1] .+ bs[i], 0f0)
    end
end

forward!(hs, Ws, bs)                        # warm up: compile kernels, create library handles
exec = instantiate(capture(() -> forward!(hs, Ws, bs)))

for batch in batches
    copyto!(hs[1], batch)                   # inputs: overwrite the captured arrays
    exec()                                  # replay all 8 operations at once
    # ... use hs[end]
end

Allocating code can be captured as well, e.g., y .= W * x .+ b, where W * x allocates a temporary: memory that is allocated during capture is allocated once, outside of the graph, kept alive by it, and overwritten by every launch.

For building graphs by hand, capture! adds captured operations to an existing graph with explicit dependencies, which makes it possible to express concurrency without juggling streams and events. update! changes the parameters of a single operation in an executable graph, e.g., to launch a kernel on different arrays without capturing again:

graph = CuGraph()
a = capture!(graph) do; x .= sin.(x); end
b = capture!(graph) do; y .= cos.(y); end
node = only(capture!(graph; after=[a; b]) do
    @cuda threads=256 combine!(z, x, y)
end)
exec = instantiate(graph)
exec()

update!(exec, node) do                      # retarget the last kernel
    @cuda threads=256 combine!(z2, x, y)
end
exec()

RTX 5080, Julia 1.13. "Submit" is the CPU time to enqueue one iteration, "per iteration" is the throughput when submitting many iterations, and "latency" includes waiting for the result:

workload		submit	per iteration	latency
MLP, 4 layers of 256, batch 16 (4 mul! + 4 broadcasts)	eager	49 µs	49 µs	55 µs
	graph	8 µs	16 µs	20 µs
RK4, 10 steps of 4096 elements (80 broadcasts, allocating)	eager	260 µs	260 µs	265–430 µs
	graph	3.5 µs	70 µs	75 µs

Eager execution is CPU-bound: the GPU sits idle between kernels, and the allocating RK4 code also suffers from GC pauses, hence its variable latency. With graphs, the GPU becomes the bottleneck: the CPU only spends half of each MLP iteration submitting work, and 5% of each RK4 iteration. The benchmark suite gains a cuda/graph group tracking eager execution, graph launch and capture of 10 small kernels.

Graphs lease the memory they use. Managed memory gets a lease count, and every way of releasing it (the pool, wrapped host memory, unsafe_free!) goes through a single release gate that postpones the release until the last lease ends. While capturing, take_ownership! records the memory that captured operations use instead of synchronizing or changing its owning stream; the resulting CuGraph, and every CuGraphExec instantiated from it, each hold a lease. Launching an executable graph takes ownership of that memory like a kernel launch does, so other tasks synchronize correctly, and the eventual release is ordered after the last launch.

Capturing no longer disables the GC. Finalizers only publish release requests (#3318). While a capture is in progress, retired resources aren't disposed of (doing so could add frees to the graph, or make prohibited API calls), and the capture drains them when it ends. Captures can't start while a disposal is in progress.

Captures don't get in the way of other tasks. Operations are captured on a dedicated non-blocking stream, which is the capturing task's stream for the duration of the capture, so other tasks can keep using the work the task submitted to its own stream earlier (e.g., an array it produced). capture(f; stream) captures on a specific stream instead. Captures use the relaxed capture mode by default: the driver checks for unsafe API calls per thread rather than per task, so in the stricter modes an unrelated task that ran on the capturing thread while the capturing task was waiting (e.g., on a lock) had its allocations or synchronizations rejected, which invalidated the capture. CUDA.jl checks the operations it performs itself; the stricter modes remain available to debug libraries that don't support capture. Loading a kernel no longer synchronizes the device while a graph is being captured, and a capture that starts while loading a kernel synchronizes the device fails instead of waiting for it.

Allocations during capture happen outside the graph. Graph-owned allocations are only valid between launches, which doesn't fit arrays whose lifetime the GC manages. Instead, memory is allocated from the regular pool on a separate stream, which owns that memory, so whoever uses it next (launching the graph, or another task) waits for the allocation like it would for any other operation. Scalars passed to libraries as CuRefs are initialized right away rather than captured, so cuBLAS calls don't add memcpy nodes, which are expensive to launch alongside kernels.

Host memory is snapshotted. A captured copy from CPU memory would read that memory on every launch, long after it may have been freed. Such copies are staged through pinned memory that the graph keeps alive, giving them snapshot semantics. Copying to CPU memory, waiting for the GPU, or using arrays from a GPUArrays.@cached scope throws a descriptive CaptureError instead of invalidating the capture. Arrays backed by HostMemory can be used to exchange data with the CPU when launching a graph.

Library allocations are handled. Libraries may allocate graph-owned memory during capture, which the driver doesn't free when the executable graph is updated or destroyed. Executable graphs track those allocations, free them when needed (like other memory, deferred while a capture is in progress), and are instantiated with AUTO_FREE_ON_LAUNCH so they can be launched repeatedly. reclaim() also trims the graph memory pool.

Limitations: random numbers generated within kernels (rand()) or by the native CUDA.RNG are the same on every launch, because their seeds are kernel arguments; rand!/CUDA.rand, which use cuRAND, are not affected. Giving captured RNG kernels a device-side seed that a tiny kernel advances at the start of every replay works, but grows KernelState (changing every kernel's ABI) and doesn't cover the GPUArrays RNG, so that's better done separately. Seeding a cuRAND generator synchronizes the device (#3336), so a task that uses cuRAND for the first time, or re-seeds it, while another task is capturing invalidates that capture. Other tasks can't wait for work submitted to an explicitly passed capture stream before the capture until it ends. Tasks spawned during capture run their operations eagerly, outside the graph. Graphs captured directly through the driver API aren't integrated with the memory management.

Tested on an RTX 5080 with Julia 1.12 (core/graph, core/resource_graph).

Fixes #3310, fixes #1230, fixes #2691. #1229 now results in a descriptive error.

[vchuravy]

What happens under task parallelism?

[codecov]

Codecov Report

❌ Patch coverage is 86.43123% with 73 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.07%. Comparing base (0bc3bde) to head (7805759).

Files with missing lines	Patch %	Lines
CUDACore/src/graph.jl	85.39%	65 Missing ⚠️
CUDACore/src/resources/backends.jl	72.72%	3 Missing ⚠️
CUDACore/src/resources/managed.jl	85.71%	3 Missing ⚠️
CUDACore/lib/cudadrv/graph.jl	88.88%	1 Missing ⚠️
CUDACore/src/array.jl	83.33%	1 Missing ⚠️

Additional details and impacted files

@@                Coverage Diff                 @@
##           tb/stream-pool    #3326      +/-   ##
==================================================
+ Coverage           84.73%   85.07%   +0.34%     
==================================================
  Files                 193      194       +1     
  Lines               19234    19674     +440     
==================================================
+ Hits                16297    16737     +440     
  Misses               2937     2937              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[maleadt]

🤖 Good question; it wasn't handled well. Concurrent captures, launching a graph from several tasks, and unrelated tasks using the GPU on other threads during a capture all worked, but:

• thread-local capture mode protects an OS thread, not a task, so tasks that ran on the capturing thread while the capture was blocked (e.g., on a lock) failed to allocate or synchronize, and invalidated the capture;
• another task loading a kernel, or using an array the capturing task had produced, got a CaptureError;
• tasks spawned inside capture executed their operations eagerly instead of capturing them.

The first two are fixed here: captures now use the relaxed mode and a dedicated stream, and kernel loading no longer synchronizes the device while capturing. Capturing the operations of spawned tasks, e.g. from libraries that parallelize internally, is #3337: they inherit the capture through a scoped value and submit to the same stream. What remains is a cuRAND issue (#3336): seeding a generator synchronizes the device, which breaks captures elsewhere.

[github-actions]

CUDA.jl Benchmarks

Details

Benchmark suite	Current: 7805759	Previous: ae637e5	Ratio
array/accumulate/Float32/1d	104780 ns	105001 ns	1.00
array/accumulate/Float32/dims=1	77078 ns	77262 ns	1.00
array/accumulate/Float32/dims=1L	1674251 ns	1684940 ns	0.99
array/accumulate/Float32/dims=2	147653 ns	148043 ns	1.00
array/accumulate/Float32/dims=2L	683993 ns	684338 ns	1.00
array/accumulate/Int64/1d	123592 ns	123503 ns	1.00
array/accumulate/Int64/dims=1	81278 ns	81902 ns	0.99
array/accumulate/Int64/dims=1L	1798607 ns	1803636 ns	1.00
array/accumulate/Int64/dims=2	162287 ns	163105 ns	0.99
array/accumulate/Int64/dims=2L	987751 ns	988275 ns	1.00
array/broadcast	17097 ns	16981 ns	1.01
array/broadcast launch	7844.75 ns	7497.25 ns	1.05
array/construct	907.5348837209302 ns	886.76 ns	1.02
array/copy	16527 ns	16325 ns	1.01
array/copyto!/cpu_to_gpu	206693 ns	206249 ns	1.00
array/copyto!/gpu_to_cpu	240638 ns	240560 ns	1.00
array/copyto!/gpu_to_gpu	8735 ns	8967 ns	0.97
array/iteration/findall/bool	139347 ns	135815 ns	1.03
array/iteration/findall/int	150731 ns	145400 ns	1.04
array/iteration/findfirst/bool	73770 ns	71077 ns	1.04
array/iteration/findfirst/int	76122 ns	73309 ns	1.04
array/iteration/findmin/1d	67663 ns	59845 ns	1.13
array/iteration/findmin/2d	103546 ns	103127 ns	1.00
array/iteration/logical	196713 ns	187405 ns	1.05
array/iteration/scalar	57398 ns	59115 ns	0.97
array/permutedims/2d	49413 ns	49364 ns	1.00
array/permutedims/3d	50813 ns	50795 ns	1.00
array/permutedims/4d	52133 ns	52271 ns	1.00
array/random/rand/Float32	10876 ns	11109 ns	0.98
array/random/rand/Int64	18272 ns	18411 ns	0.99
array/random/rand!/Float32	8205 ns	8183 ns	1.00
array/random/rand!/Int64	15968 ns	15765 ns	1.01
array/random/randn/Float32	34508 ns	34501 ns	1.00
array/random/randn!/Float32	24482 ns	24403 ns	1.00
array/reductions/mapreduce/Float32/1d	32756 ns	30573 ns	1.07
array/reductions/mapreduce/Float32/dims=1	39733 ns	40213 ns	0.99
array/reductions/mapreduce/Float32/dims=1L	53453 ns	53356 ns	1.00
array/reductions/mapreduce/Float32/dims=2	59099 ns	59358 ns	1.00
array/reductions/mapreduce/Float32/dims=2L	71457 ns	71210 ns	1.00
array/reductions/mapreduce/Int64/1d	39636 ns	37469 ns	1.06
array/reductions/mapreduce/Int64/dims=1	42923 ns	42915 ns	1.00
array/reductions/mapreduce/Int64/dims=1L	89248 ns	89249 ns	1.00
array/reductions/mapreduce/Int64/dims=2	60318 ns	61370 ns	0.98
array/reductions/mapreduce/Int64/dims=2L	87612 ns	87543 ns	1.00
array/reductions/reduce/Float32/1d	33070 ns	30615 ns	1.08
array/reductions/reduce/Float32/dims=1	40252 ns	40264 ns	1.00
array/reductions/reduce/Float32/dims=1L	52780 ns	52907 ns	1.00
array/reductions/reduce/Float32/dims=2	59232 ns	59462 ns	1.00
array/reductions/reduce/Float32/dims=2L	71502 ns	71728 ns	1.00
array/reductions/reduce/Int64/1d	39814 ns	37570 ns	1.06
array/reductions/reduce/Int64/dims=1	42898 ns	42862 ns	1.00
array/reductions/reduce/Int64/dims=1L	89291 ns	89319 ns	1.00
array/reductions/reduce/Int64/dims=2	60379 ns	61437 ns	0.98
array/reductions/reduce/Int64/dims=2L	87505 ns	86880 ns	1.01
array/reverse/1d	16342 ns	16128 ns	1.01
array/reverse/1dL	70311 ns	70115 ns	1.00
array/reverse/1dL_inplace	67738 ns	67602 ns	1.00
array/reverse/1d_inplace	9435 ns	9293 ns	1.02
array/reverse/2d	20934 ns	21118 ns	0.99
array/reverse/2dL	74766 ns	74528 ns	1.00
array/reverse/2dL_inplace	67727 ns	67778 ns	1.00
array/reverse/2d_inplace	10840 ns	10676 ns	1.02
array/sorting/1d	2846913 ns	2846981 ns	1.00
array/sorting/2d	1089420 ns	1088944 ns	1.00
array/sorting/by	3400808 ns	3400388 ns	1.00
cuda/graph/capture	28517 ns
cuda/graph/eager	40170 ns
cuda/graph/launch	14280 ns
cuda/synchronization/context/auto	7793.75 ns	6766.6 ns	1.15
cuda/synchronization/context/blocking	867.6304347826087 ns	793.7692307692307 ns	1.09
cuda/synchronization/context/nonblocking	7888.5 ns	6805.4 ns	1.16
cuda/synchronization/stream/auto	721.5116279069767 ns	691.0939597315436 ns	1.04
cuda/synchronization/stream/blocking	950.9545454545455 ns	868.875 ns	1.09
cuda/synchronization/stream/nonblocking	8225.666666666666 ns	7338.5 ns	1.12
integration/byval/reference	148090 ns	147882 ns	1.00
integration/byval/slices=1	149454 ns	149033 ns	1.00
integration/byval/slices=2	291606 ns	291277 ns	1.00
integration/byval/slices=3	434100 ns	434160 ns	1.00
integration/cudadevrt	105565 ns	105116 ns	1.00
integration/volumerhs	9207466 ns	9206594 ns	1.00
kernel/indexing	11847 ns	11855 ns	1.00
kernel/indexing_checked	12296 ns	12209 ns	1.01
kernel/launch	2498.1111111111113 ns	2433 ns	1.03
kernel/occupancy	927.6129032258065 ns	1008.8235294117648 ns	0.92
kernel/rand	15137 ns	14899 ns	1.02
latency/import	4335736926 ns	4305717656 ns	1.01
latency/precompile	5091803797 ns	5095420673 ns	1.00
latency/ttfp	4790048481 ns	4793974717 ns	1.00

This comment was automatically generated by workflow using github-action-benchmark.

[AntonOresten]

Super stoked about this! Nice work.

🤖 A performance observation: launching an executable graph costs about 70–80 ns of CPU time per array the graph holds, on top of the driver launch.

Each exec() goes through with_ordered_managed(exec.launch_memory) (graph.jl#L824), which locks and takes ownership of every leased array, even when the graph is relaunched on the same stream and nothing else has touched those arrays since. The benchmarks in the PR hold at most about a hundred arrays (RK4), where this adds only a few µs, but a captured training step (parameters, gradients, optimizer state) easily holds thousands:

Click to see setup

using CUDA

# CPU time of `f()`, with the GPU idle before each call
function cputime(f; reps=300)
    minimum(1:5) do _
        sum(1:reps) do _
            synchronize(); @elapsed f()
        end / reps
    end * 1e6
end

for N in (10, 100, 1000, 3000)
    xs = [CUDA.zeros(Float32, 16) for _ in 1:N]
    exec = instantiate(capture(() -> foreach(x -> x .+= 1, xs)))
    exec(); synchronize()
    s = stream()
    t_exec = cputime(() -> exec())
    t_raw  = cputime(() -> CUDA.CUDACore.cuGraphLaunch(exec, s))
    println("$N arrays: exec() $(round(t_exec, digits=1)) µs, cuGraphLaunch $(round(t_raw, digits=1)) µs")
end

On a GB10 (tb/graphs at 7805759), the GPU idle before each call:

10 arrays:   exec() 2.2 µs,   cuGraphLaunch 1.6 µs
100 arrays:  exec() 6.1 µs,   cuGraphLaunch 1.6 µs
1000 arrays: exec() 85.8 µs,  cuGraphLaunch 15.1 µs
3000 arrays: exec() 243.3 µs, cuGraphLaunch 19.6 µs

With the rest of the stack on top it's higher: about 110–130 µs at 1000 arrays and 350–575 µs at 3000 (measured at #3328's head, 38d4a78, and at #3360, a581505), so a fast path would need to cover what the later PRs add per array too.

It's still far cheaper than launching those kernels without a graph, but it's 6–12× the raw launch, and grows with the size of the model. Would a fast path make sense, one that skips the loop when the graph is relaunched on the same stream and no other task has used its memory since the previous launch?

[AntonOresten]

🤖 A second observation, about memory: temporaries allocated during capture are leased for the graph's lifetime and never reused within it, so a graph holds the sum of all its temporaries. Code that runs fine eagerly can then fail to capture:

Click to see setup

using CUDA

# 50 allocating steps; each temporary is dead as soon as the next one exists
function step(x)
    y = x
    for _ in 1:50
        y = y .* 0.99f0 .+ 0.01f0
    end
    return y
end

# the same, but freeing each temporary explicitly once it's dead
function step_free(x)
    y = x
    for i in 1:50
        z = y .* 0.99f0 .+ 0.01f0
        i > 1 && CUDA.unsafe_free!(y)
        y = z
    end
    return y
end

x = CUDA.rand(Float32, 4 * 2^20)         # 16 MiB; 50 temporaries = 800 MiB

for _ in 1:10; step(x); end; synchronize()
println("eager, 10 times: ok")

for f in (step, step_free)
    try
        g = capture(() -> f(x))
        println("capture $f: ok")
    catch err
        println("capture $f: ", first(sprint(showerror, err), 80))
    end
    GC.gc(true); CUDA.reclaim()
end

With JULIA_CUDA_HARD_MEMORY_LIMIT=400MiB on a GB10 (tb/graphs at 7805759, and the same at #3360):

eager, 10 times: ok
capture step: Out of GPU memory trying to allocate 16.000 MiB
capture step_free: Out of GPU memory trying to allocate 16.000 MiB

Explicitly freeing the temporaries with CUDA.unsafe_free! during the capture doesn't help either, so even code that knows when its temporaries die can't reduce this today.

The docs mention this and recommend capturing in-place code for large problems. But making allocating code capturable is what makes this PR attractive for e.g. training steps. PyTorch reuses freed blocks within a capture, though that's easier with refcounting than with a GC. Is reuse within a graph something you'd consider, here or as a follow-up?

[maleadt]

@AntonOresten Please see #3362