CUDA Stream Capture, in depth: recording stream work into graphs, fork and join, capture modes, invalidation, memory and graph updates

What capture records, and what it freezes

During capture, a stream stops being a queue for the GPU and becomes a recorder. Each kernel launch, cudaMemcpyAsync, cudaMemsetAsync, host function launch and event record or wait issued to a capturing stream becomes a node, and the stream's ordering becomes edges. Nothing runs. When you call cudaStreamEndCapture you get a cudaGraph_t: a template. cudaGraphInstantiate then validates it and does the setup work once, producing a cudaGraphExec_t you launch with cudaGraphLaunch.

The consequence that causes most bugs: everything the CPU computed while capturing is baked in. Kernel arguments are copied by value at capture time, so a device pointer, a sequence length or a scalar learning rate passed as an argument is a constant in the graph. Grid and block sizes are constants. If your host code takes a branch based on a value, only the branch taken during capture exists in the graph, and host code between launches does not run on replay unless captured with cudaLaunchHostFunc. A graph replays the same work on the same addresses: change inputs by writing into the captured buffers.

The capture API in one program

The program below captures a step that forks onto a second stream and joins back. step is ordinary stream code. Run it once before capture: first calls may trigger lazy module loading, and libraries such as cuBLAS may allocate workspaces on first use, which must not happen inside a capture.

// CUDA 12.x. CK() aborts with file, line and cudaGetErrorString on any error.
void step(cudaStream_t s0, cudaStream_t s1, cudaEvent_t fork, cudaEvent_t join) {
  CK(cudaEventRecord(fork, s0));
  CK(cudaStreamWaitEvent(s1, fork, 0));     // s1 joins the capture here
  proj<<<grid, block, 0, s0>>>(d_in, d_q);
  side<<<grid, block, 0, s1>>>(d_in, d_k);  // no edge to proj: may run concurrently
  CK(cudaEventRecord(join, s1));
  CK(cudaStreamWaitEvent(s0, join, 0));     // rejoin before EndCapture
  combine<<<grid, block, 0, s0>>>(d_q, d_k, d_out);
}

int main() {
  cudaStream_t s0, s1;
  CK(cudaStreamCreateWithFlags(&s0, cudaStreamNonBlocking));
  CK(cudaStreamCreateWithFlags(&s1, cudaStreamNonBlocking));
  cudaEvent_t fork, join;
  CK(cudaEventCreateWithFlags(&fork, cudaEventDisableTiming));
  CK(cudaEventCreateWithFlags(&join, cudaEventDisableTiming));

  step(s0, s1, fork, join);                 // warm-up run: lazy loading, workspaces
  CK(cudaDeviceSynchronize());

  cudaGraph_t graph;
  CK(cudaStreamBeginCapture(s0, cudaStreamCaptureModeGlobal));
  step(s0, s1, fork, join);                 // same calls, recorded instead of run
  CK(cudaStreamEndCapture(s0, &graph));

  cudaGraphExec_t exec;
  CK(cudaGraphInstantiate(&exec, graph, 0));
  for (int it = 0; it < iters; ++it) {
    CK(cudaMemcpyAsync(d_in, h_pinned[it], bytes, cudaMemcpyHostToDevice, s0));
    CK(cudaGraphLaunch(exec, s0));          // one submission for all three kernels
  }
  CK(cudaStreamSynchronize(s0));
  CK(cudaGraphExecDestroy(exec));
  CK(cudaGraphDestroy(graph));
}

Notice the input refresh. The copy into d_in is issued outside the graph, on the same stream, before each launch, so stream order guarantees the graph sees the new data. The graph itself always reads from the same d_in address that existed during capture.

Several streams, one capture

Capture starts on one origin stream. Another stream joins the capture the moment it waits on an event recorded in a capturing stream; the programming guide describes this as the wait placing the stream in capture mode, with its next item depending on the nodes in the captured event. From then on, work on s1 becomes nodes in the same graph, and the event edges become graph edges. Because the fork is recorded before proj, the graph has proj and side as independent roots feeding combine, and on replay the GPU is free to run them concurrently. Had the fork been recorded after proj, side would depend on it and the graph would run serially: event placement is the dependency.

Every stream that joined must rejoin the origin stream before cudaStreamEndCapture, or ending the capture fails with cudaErrorStreamCaptureUnjoined. The usual culprit is a helper that forks for an asynchronous copy and never joins, which was harmless outside capture.

Capture modes: what they actually restrict

The mode argument of cudaStreamBeginCapture is widely misunderstood. It does not change what gets recorded. It controls which threads are barred from calls the runtime considers potentially unsafe during capture, the main example being cudaMalloc, which is not stream-ordered and so cannot become part of the graph. The runtime reference defines them this way:

• Global (the default): a thread may not make unsafe calls if it has its own ongoing capture that is not relaxed, or if any other thread has a concurrent capture begun in global mode.
• ThreadLocal: a thread is restricted only by its own ongoing non-relaxed capture. Other threads may keep allocating.
• Relaxed: the capturing thread itself is not prohibited from unsafe calls. You take responsibility for not, for example, freeing memory a captured node will use.

Global mode turns a data-loader thread's cudaMalloc during your capture into an error instead of a silent hazard; in multi-threaded servers that is often too strict, hence ThreadLocal. A thread can also change its own interaction mode temporarily with cudaThreadExchangeStreamCaptureMode, which is how libraries make a known-safe call during someone else's capture. PyTorch exposes the same three choices as capture_error_mode="global", "thread_local" and "relaxed".

What breaks a capture

Some operations are simply invalid while a stream is capturing, because captured work has not been scheduled and so has no status to report. The guide states that synchronizing or querying a capturing stream or a captured event is invalid, and that if a stream was not created with cudaStreamNonBlocking, using the legacy default stream during capture is invalid too. An invalid call does more than fail: it invalidates the capture, so every later call on the capturing streams also errors until you call cudaStreamEndCapture, which then returns no usable graph.

To check whether you are inside a capture, for instance in a library function that would otherwise synchronize, call cudaStreamIsCapturing. cudaStreamGetCaptureInfo returns more: status, a capture id, the graph being built and the current dependency set. Its signature has changed across CUDA 11, 12 and 13, so check the header for your toolkit before calling it.

Memory during capture

Since cudaMalloc cannot be captured, the classic approach is to allocate every buffer before capture and keep it alive as long as the executable graph exists. The stream-ordered allocator changes this: cudaMallocAsync and cudaFreeAsync issued to a capturing stream become memory allocation and free nodes. If the graph allocates without freeing, the memory stays owned by the graph after launch, and relaunching it is an error unless you free it first or instantiate with cudaGraphInstantiateFlagAutoFreeOnLaunch.

PyTorch takes a third route. Its caching allocator gives each capture a private memory pool, so tensors created inside the captured region keep fixed addresses for the life of the graph. Two graphs can share one pool by passing the same pool token, from torch.cuda.graph_pool_handle() or other_graph.pool(). Sharing is safe when you replay the graphs in the order you captured them and never run them concurrently. It is the standard trick for keeping memory flat when you capture many shapes.

Updating instead of re-instantiating

Instantiation is the expensive part, so when only parameters change, such as a pointer or a scalar, recapture into a new template and apply it to the existing executable graph with cudaGraphExecUpdate. The new graph must have identical topology and dependency order, so your capturing code must issue the same calls in the same order. When the update is refused, fall back to a fresh instantiate.

// Same topology, new parameters: try an in-place update before re-instantiating.
cudaGraph_t g2;
CK(cudaStreamBeginCapture(s0, cudaStreamCaptureModeThreadLocal));
step_with(s0, s1, fork, join, d_in_b, d_out_b);   // identical call order, new pointers
CK(cudaStreamEndCapture(s0, &g2));

cudaGraphExecUpdateResultInfo info;
if (cudaGraphExecUpdate(exec, g2, &info) != cudaSuccess) {
  (void)cudaGetLastError();                        // clear the non-sticky error
  CK(cudaGraphExecDestroy(exec));                  // topology or node type changed
  CK(cudaGraphInstantiate(&exec, g2, 0));
}
CK(cudaGraphDestroy(g2));

Recapturing and updating is cheaper than instantiating but not free, so a workload whose shapes change every step gains little from graphs.

Capture in PyTorch

In PyTorch, torch.cuda.graph is a context manager around capture of the current stream; the documentation marks it as beta. The pattern below is the one the PyTorch docs teach: warm up on a side stream, capture with static input tensors, then copy new data into those tensors and call replay().

import torch

model = Net().cuda().eval()
static_x = torch.zeros(32, 4096, device="cuda", dtype=torch.float16)

# Warm up on a side stream so lazy init and autotuning happen before capture.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s), torch.no_grad():
    for _ in range(3):
        model(static_x)
torch.cuda.current_stream().wait_stream(s)

g = torch.cuda.CUDAGraph()
with torch.no_grad(), torch.cuda.graph(g, capture_error_mode="thread_local"):
    static_y = model(static_x)          # recorded, addresses now fixed

def run(x):
    static_x.copy_(x)                   # new data into the captured input buffer
    g.replay()                          # same kernels, same addresses
    return static_y.clone()             # static_y is overwritten by the next replay

The .clone() is easy to forget: static_y lives in the graph's pool and the next replay overwrites it. For training, torch.cuda.make_graphed_callables captures forward and backward of a module, and torch.compile(mode="reduce-overhead") applies CUDA graphs automatically; see torch.compile in depth.

Worked example: bucketed LLM decode

Take an LLM server whose decode step runs one forward pass for a batch of sequences, one token each. Batch size changes every step as requests arrive and finish, but a graph has fixed shapes. The usual design is bucketing: capture one graph per batch size in a set such as 1, 2, 4, 8, 16 and 32, pad each live batch up to the next bucket, and replay that bucket's graph. Larger live batches fall back to eager execution.

The data flow per step is: the scheduler picks the bucket; the host writes token ids, positions and the block table for the paged KV cache into that bucket's static input tensors with asynchronous copies; the graph replays; sampling reads the static logits. Everything that varies, including each sequence's length, must arrive as data in device memory, never as a Python integer passed to a kernel, because the integer would be frozen.

Six graphs with private pools would hold six sets of activations. Capturing them into one shared pool, largest bucket first, lets smaller graphs reuse the largest one's blocks; the shared-pool rule holds because the server replays one bucket at a time on one stream. Verify in an Nsight Systems timeline: the gaps between decode kernels should close and CPU time per step should fall.

Failure modes

• Stale scalars. A step counter or learning rate passed as a kernel argument is frozen. Keep it in a device tensor and update it in place.
• Reallocated inputs. Assigning a new tensor to the Python variable instead of copy_ into the static one means the graph reads old data, with no error.
• Hidden syncs. A shape check, a .item(), a CPU-side assertion or a library that queries a stream invalidates the capture. Find it with cudaStreamIsCapturing guards or by bisecting the captured region.
• Memory blow-up. Many shapes, each with a private pool, can exhaust memory. Share pools or capture fewer buckets.
• Randomness. RNG offsets must advance on replay; PyTorch registers its generators with the graph, but custom RNG state passed by value repeats the same numbers every replay.
• Update rejected. A conditional branch that captures a different node sequence makes cudaGraphExecUpdate fail. Keep capturing code branch-free or re-instantiate.

Trade-offs and alternatives

Graphs trade flexibility for submission cost. You gain fewer CPU launches and less jitter between kernels; you pay with static shapes and addresses, instantiation time, memory per captured shape and harder debugging. They help most when kernels are small and numerous, as in decode, and least when single kernels run for milliseconds. Alternatives: kernel fusion cuts the number of launches, and overlapping copies with compute, as in asynchronous memcpy, hides latency rather than removing it. Many production stacks use all three.

What to do next

1. Profile one step and count kernels and launch gaps before deciding graphs will help.
2. Make every input a pre-allocated device buffer and every varying scalar a device tensor.
3. Run the step once outside capture to warm up libraries, then capture it.
4. Check that every forked stream records a join event that the origin stream waits on.
5. Pick ThreadLocal mode in multi-threaded processes unless you need global protection.
6. Validate replay against eager output on several inputs, including a changed input.
7. Bucket shapes and share one memory pool, capturing the largest bucket first.
8. Use cudaGraphExecUpdate for parameter changes and fall back to re-instantiate.
9. Re-profile and keep an eager fallback path behind a flag for debugging.