CUDA Graph

A CUDA Graph is constructed in two distinct phases. First, a capture sequence is initiated, during which all subsequent CUDA API calls (e.g., kernel<<<>>>(), cudaMemcpyAsync) are recorded into a graph object instead of executing immediately. This creates a template of operations. The graph is then instantiated into an executable graph, which is a lightweight, optimized representation ready for repeated replay. This separation allows for one-time setup costs to be amortized over many executions.

The primary performance benefit stems from eliminating per-launch overhead. In a traditional CUDA stream, each kernel launch incurs costs for:

• Driver and runtime latency
• Scheduling work onto the GPU
• Kernel argument marshaling

A CUDA Graph bundles all work into a single, coarse-grained operation. When the graph is launched (cudaGraphLaunch), the entire sequence is submitted with the overhead of a single launch, not N launches. This is particularly impactful for workloads comprising many small, rapid-fire kernels.

A CUDA Graph represents a static workflow. The graph's structure—the kernels, their dependencies, and the memory addresses they operate on—is fixed at instantiation. This allows the CUDA driver to perform aggressive optimizations upfront:

• Optimal node scheduling and dependency resolution.
• Pre-allocation of resources like GPU registers and shared memory.
• Pre-fetching of kernel arguments to GPU-accessible memory.

This static nature means graphs are ideal for inference loops, solvers, and other compute patterns where the operation sequence is invariant across iterations, even if the input data changes.

While the graph structure is static, execution parameters can be updated between launches without re-capturing. This is achieved through:

• Graph node parameters: Kernel parameters (e.g., pointers to input/output buffers) can be updated using cudaGraphExecKernelNodeSetParams.
• Conditional nodes: Introduced in CUDA 10, child graphs and conditional nodes (e.g., cudaGraphAddChildGraphNode, cudaGraphAddConditionalNode) allow for limited dynamism, where sub-graphs can be executed based on a runtime condition, blending static optimization with runtime flexibility.

Major frameworks leverage CUDA Graphs to accelerate inference. For example:

• TensorRT uses graphs to capture the entire engine execution.
• PyTorch's torch.cuda.CUDAGraph API allows users to capture model-forward passes.
• Triton Inference Server employs graphs to capture the end-to-end pre-process/model/post-process pipeline.

The typical pattern is warm-up, capture, then replay. The first iteration runs normally to allocate memory and populate caches; this execution is captured into a graph; all subsequent inferences replay the efficient graph.

CUDA Graphs and kernel fusion are complementary but distinct optimization layers. Kernel fusion is a compile-time optimization that merges multiple computational operators (e.g., Conv + BatchNorm + ReLU) into a single, custom fused kernel to reduce global memory traffic. CUDA Graph is a launch-time optimization that groups multiple, potentially independent, kernel launches (fused or not) into a single submission unit to reduce CPU-driven launch latency. An optimized pipeline often uses both: fused kernels within a CUDA Graph.

Kernel fusion is a low-level compiler optimization that combines multiple, separate GPU computational kernels into a single, unified kernel. This eliminates the kernel launch overhead and intermediate memory stores/loads between each operation, improving data locality and overall throughput. It is a core technique used within the kernels that a CUDA Graph captures and replays.

• Primary Goal: Reduce launch latency and improve memory bandwidth utilization.
• Implementation Level: Hand-written or compiler-generated at the CUDA C/C++ kernel level.
• Example: Fusing an element-wise addition, ReLU activation, and scaling operation into one kernel.

Operator fusion is a graph-level optimization performed by deep learning compilers (e.g., XLA, TVM, MLIR) that merges adjacent computational nodes in a neural network's computational graph. The fused subgraph is then typically implemented via a fused kernel. While CUDA Graphs capture a sequence of already-launched operations, operator fusion creates new, more efficient operations to be launched.

• Primary Goal: Minimize intermediate tensor memory transfers between operators.
• Scope: Works on the abstract computational graph before kernel code is generated.
• Canonical Pattern: The Fused Conv-BN-ReLU, which combines Convolution, Batch Normalization, and activation.

Kernel launch overhead is the fixed latency and resource cost incurred each time the CPU instructs the GPU to execute a kernel. This includes API call processing, driver scheduling, and setting up kernel arguments. For workloads with many small, sequential kernels, this overhead can dominate execution time.

• CUDA Graph's Role: A CUDA Graph amortizes this cost by capturing the entire workflow (kernels & memory ops) and replaying it with a single, low-overhead launch.
• Fusion's Role: Kernel and operator fusion directly reduce the number of launches required.
• Impact: Most significant for high-throughput inference serving and models with many small operations.

Graph fusion is the automated process within a compiler where a fusion planner identifies profitable subgraphs (or fusion groups) to merge. It uses fusion heuristics and a cost model for fusion to analyze fusion profitability based on data dependencies, memory access patterns, and operator types. This is the algorithmic backbone that enables operator fusion.

• Techniques: Includes vertical fusion (producer-consumer chains) and horizontal fusion (parallel operators).
• Compiler Integration: Core component of XLA, TVM, and MLIR compilation pipelines.
• Relationship to CUDA Graph: CUDA Graph is a runtime, coarse-grained capture; graph fusion is a compile-time, fine-grained optimization that often produces the kernels a graph will capture.

Fused Multi-Head Attention is a premier example of an extremely profitable fused kernel critical for transformer models. It implements the entire attention mechanism—query/key/value projection, scoring, softmax, and aggregation—in a single, highly optimized GPU kernel. This avoids materializing the large intermediate attention matrix to global memory.

• Breakthrough Implementation: FlashAttention is an IO-aware, fused attention algorithm that further optimizes by keeping data in SRAM and recomputing on-chip, achieving unprecedented speed and memory efficiency for long sequences.
• Performance Impact: Can improve attention throughput by an order of magnitude and reduce memory footprint.
• Context: Such a fused kernel is a perfect candidate for capture and efficient replay within a CUDA Graph.

Just-In-Time (JIT) Fusion is a runtime optimization where the decision of which operators to fuse and the generation of the corresponding fused kernel occurs dynamically during model execution. This allows for specialization based on actual input shapes and runtime context.

• Framework Example: PyTorch's torch.compile uses a JIT compiler (Inductor) to capture a model's graph and apply fusion optimizations just before execution.
• Advantage: Flexibility and adaptability to dynamic graph structures.
• Contrast with AOT: Differs from Ahead-of-Time (AOT) Fusion, where fusion is fixed during static compilation. CUDA Graph capture is a form of JIT optimization at the launch level, often working on a graph already optimized by JIT fusion.