Graph Capture and CUDA Graphs

Every time PyTorch dispatches a kernel to the GPU, it pays a tax: Python overhead, C++ dispatch, and a CUDA API call that wakes the driver. For a transformer forward pass with hundreds of small kernels, this CPU-side chatter can cost more wall-clock time than the actual arithmetic. On an A100 running BERT at batch size 1, CPU overhead can account for 30-50% of total latency. CUDA Graphs are the mechanism that cuts that tax to near zero.

What the CUDA Graph Model Actually Is

A CUDA Graph is a directed acyclic graph (DAG) where each node represents a GPU operation: a kernel launch, a memory copy, a memory set, or a host-function call. Edges encode dependencies. Instead of submitting nodes one at a time through the CUDA stream, you submit the entire DAG in a single cudaGraphLaunch() call.

The workflow has three phases:

1. Capture. Wrap your workload between cudaStreamBeginCapture() and cudaStreamEndCapture(). During capture the driver records every CUDA API call but does not execute them. You end up with a cudaGraph_t object representing the DAG.
2. Instantiation. Call cudaGraphInstantiate() to compile the DAG into an executable form (cudaGraphExec_t). This is a one-time cost, typically a few hundred microseconds.
3. Replay. Call cudaGraphLaunch(exec, stream) as many times as you like. Each call launches the full DAG with one driver interaction instead of hundreds.

The GPU still executes the same kernels in the same order. What changes is the cost of telling it to. NVIDIA's own measurements put individual kernel launch overhead at roughly 2-4 microseconds; for a 300-kernel forward pass this adds up to ~1 ms of pure overhead per iteration, which CUDA Graphs eliminates.

PyTorch's Graph API

PyTorch exposes CUDA Graphs through two interfaces:

torch.cuda.CUDAGraph (low-level):

g = torch.cuda.CUDAGraph()

# Warmup: run the model once on the real stream so caches are hot
with torch.cuda.stream(torch.cuda.Stream()):
    for _ in range(3):
        y = model(x)

# Capture
with torch.cuda.graph(g):
    y = model(x)

# Replay (x and y are now fixed memory addresses)
x.copy_(new_input)
g.replay()
result = y.clone()

The critical point is that x and y are static buffers. The graph captures the tensor addresses, not their contents. To feed new inputs you copy data into the same allocation; to read output you copy from the same allocation. This single constraint drives most of the practical complexity.

torch.cuda.make_graphed_callables (module-level):

Wraps individual nn.Module or functions and returns graph-accelerated versions, handling the static-buffer bookkeeping automatically. This is safer for models where only part of the computation is graph-safe.

torch.compile with mode="reduce-overhead":

The highest-level option. torch.compile invokes TorchDynamo to trace the graph and TorchInductor to generate kernels, then wraps captured regions in CUDA Graphs automatically. The reduce-overhead mode is the intended path for most users since PyTorch 2.0.

How the Capture-Replay Contract Works

Replay is only valid if the graph structure is identical on each invocation. This means:

• All tensor shapes must be the same (no dynamic sequence lengths without padding to a fixed length).
• No CPU-GPU synchronisation inside the graph (.item(), .numpy(), print() of a GPU tensor, in-place ops that branch on tensor values).
• No Python control flow that changes the set of kernels launched (no if loss > threshold: do_something()).
• Allocations inside the graph must be to the same addresses each time; PyTorch achieves this with a memory pool that resets to its original state between replays.

The memory pool rule has a subtle implication: intermediate activations allocated during the captured forward pass are not freed between replays. The pool grows to its peak usage during capture and stays there. For a large model this can add significant memory pressure on top of the model's normal footprint.

CUDA Graphs also do not support RNG operations by default unless you use torch.cuda.make_graphed_callables with explicit RNG state management, because random seeds are stateful and vary between launches.

The cudaGraphExecUpdate Escape Hatch

Instantiation is expensive enough (~0.5-1 ms) that re-instantiating on every shape change is impractical. CUDA 10.2+ provides cudaGraphExecUpdate(), which patches an existing cudaGraphExec_t in-place when only kernel parameters (not graph topology) change. PyTorch uses this under the hood in torch.compile's CUDA graph trees implementation to handle weight updates between training steps without full re-instantiation.

The distinction matters for training: you typically run the same graph for hundreds of batches. After an optimiser step the parameter tensors hold new values, but because those tensors are the same allocations the graph already refers to, the next replay() automatically uses the updated weights. No update call needed. The static-memory-address property is an asset here, not just a constraint.

torch.compile and CUDA Graph Trees

torch.compile(model, mode="reduce-overhead") uses a structure called CUDA graph trees, introduced to handle the complication that a full training loop contains branchy Python code (logging, gradient clipping, conditional early stopping) that cannot all be captured in one graph.

The implementation breaks the computation into segments, each of which is individually capture-safe, then chains them into a tree. On the first few invocations it runs eagerly to identify which paths are taken; then it captures and caches each hot path. When a new shape or path is encountered it falls back to eager execution and optionally re-captures. This adaptive strategy means you get graph acceleration on the hot path without sacrificing correctness on the cold path.

A rough performance picture from the PyTorch 2.0 benchmarks: reduce-overhead mode yields 20-40% latency reduction on small-to-medium models where CPU overhead was the bottleneck, and negligible gain on large models where GPU compute dominates.

When It Falls Down

Dynamic shapes. Any model that processes variable-length sequences without padding to a fixed maximum will cause graph invalidation on nearly every call, making capture overhead a net loss. Workarounds: bucket inputs into a small set of fixed lengths (e.g., powers of two), or use torch.compile(dynamic=True) which attempts symbolic shape reasoning but cannot always produce a capturable graph.

In-graph synchronisation. Loss scaling in mixed-precision training often checks for inf/nan values using torch.cuda.amp.GradScaler, which calls .item() internally. This forces a CPU-GPU sync that poisons the capture. Solution: structure the inf check to run outside the captured region, or use torch.compile which handles this automatically.

Memory overhead. The graph memory pool holds peak activations alive for the full training run. On a 70B-parameter model or a very deep architecture with large intermediates, this can be prohibitive. Monitor with torch.cuda.memory_summary().

Graph re-instantiation cost. If shapes change frequently (variable-length text, images of different resolutions), the amortised instantiation cost can exceed the savings from eliminating kernel launch overhead. Profile before committing.

Debugging. Errors inside a captured graph surface at replay time with minimal context, because the actual execution is deferred. Standard CUDA error checking via cuda-memcheck or Compute Sanitizer still works but requires special flags to handle the deferred execution model.

Multi-GPU (DDP/FSDP). Collective operations (AllReduce, AllGather) can be captured, but only if all ranks enter and exit the capture window simultaneously. Any asymmetric CPU logic between ranks causes a deadlock. As a result, DDP gradient communication is often excluded from the captured region and runs eagerly.

Further Reading

• Accelerating PyTorch with CUDA Graphs (PyTorch Blog, 2021) - the canonical introduction with concrete benchmarks on Mask R-CNN and BERT.
• CUDA Graphs (NVIDIA Developer Blog, 2019) - NVIDIA's original announcement; covers the stream-capture API with C++ examples and microsecond-level overhead measurements.
• CUDA C++ Programming Guide (NVIDIA) - the authoritative reference for the underlying cudaGraph_t API, node types, and cudaGraphExecUpdate.