Parameterized CUDA Graph Launch in PyTorch: CUDA Graphs Without the Pain
summary
Explains a new Parameterized CUDA Graph launch that lets CUDA Graphs replay with new tensor pointers by patching kernel pointer args and borrowing memory from the default allocator instead of a private pool. Works with torch.compile and reduces partial-capture memory overhead — llama3.1-70B in TorchTitan shows a 2 GB memory increase versus ~10 GB with traditional graphs. Implements pointer detection via double-capture, runs on TorchBench single-GPU, and still trips on distributed, NV SHMEM and alignment edge cases.
AI-written summary · about
chapters (10)
- 0:00CUDA graphs in PyTorch
- 1:02What are CUDA graphs
- 3:10Performance gains
- 6:40Current usage
- 9:49Challenges today
- 12:57Parameterized CUDA graph launch
- 15:03Solving challenges
- 17:32Why implementation is hard
- 22:48Alternative techniques
- 23:51Conclusion
related talks
12mLightning Talk: CUDAGraph in a Partial Graph World
43mExploring Deep Learning Framework PyTorch
25mLifecycle of a Parameter
19mNeuGraph: Parallel Deep Neural Network Computation on Large Graphs
15mIdentifying and Analyzing Pitfalls in GNN Systems
36mPyTorch: Fast Differentiable Dynamic Graphs in Python
18mTC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs
15mGraCE: Unlocking CUDA Graphs with Compiler Support for ML Workloads