Shape-Stable Dynamic Control Flow in PyTorch CUDA Graphs

CUDA Graphs eliminate CPU overhead and are critical for maximum performance in deep learning today. However, they have historically required execution to be static: fixed launch topology, fixed tensor shapes, and no data-dependent decisions that would force a CPU roundtrip.

This talk presents a new contract for CUDA Graphs in PyTorch. Many workloads are dynamic in a structured way: so long as the control-flow region has a static output shape, CUDA Graphs can use a feature called conditional nodes to support dynamically computed branch choices and loop guards, as well as a bounded number of dynamic input shapes. The control flow ops torch.cond(), torch.while_loop(), and torch.switch() encode this contract precisely. If a workload can be expressed using these operators, PyTorch will lower the control flow to IF, WHILE, and SWITCH conditional nodes during graph capture.

We will demonstrate how this unlocks CUDA graphs for important workloads that were previously incompatible, such as skipping optimizer steps with non-finite gradients, MoE routing with fixed-capacity dispatch buffers, embedding-bag reductions for recsys, and numerical solvers that run until convergence.