Parameterized CUDA Graph Launch in PyTorch: CUDA Graphs Without the Pain - Daniel Galvez, NVIDIA

This talk focuses on the challenges of using CUDA Graphs in PyTorch, highlighting the potential bottleneck caused by CPU kernel launch overhead on modern GPUs. The speaker discusses common failures in CUDA Graph capture that PyTorch users encounter and provides guidance on how to diagnose and fix these issues. Additionally, the session introduces an innovative solution: Parameterized CUDA Graph launch, which applies CUDA Graphs only to compatible regions of a workload, significantly simplifying user efforts. Leveraging this approach, workflows such as llama3.1-70B in torchtitan can achieve substantial memory efficiency, showing only a 2 GB increase over a non-graph baseline, in contrast to more traditional methods that may require up to 10 GB.