# PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch

arXiv: https://arxiv.org/abs/2503.19779 · HTML: https://mlsanity.com/paper/2503.19779/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 90 references that are in the map.
- First posted on arXiv 25 Mar 2025.

## Builds on (most cited first, 23 of 23)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://mlsanity.com/paper/2403.02310/) (2024)
- [Caffe: Convolutional Architecture for Fast Feature Embedding](https://mlsanity.com/paper/1408.5093/) (2014)
- [InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management](https://mlsanity.com/paper/2406.19707/) (2024)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Parrot: Efficient Serving of LLM-based Applications with Semantic Variable](https://mlsanity.com/paper/2405.19888/) (2024)
- [Analysis of Large-Scale Multi-Tenant GPU Clusters for DNN Training Workloads](https://mlsanity.com/paper/1901.05758/) (2019)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference](https://mlsanity.com/paper/2410.18038/) (2024)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inference](https://mlsanity.com/paper/2006.03031/) (2020)
- [Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performance](https://mlsanity.com/paper/2110.15238/) (2021)
- [The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding](https://mlsanity.com/paper/2110.10221/) (2021)
- [LazyTensor: combining eager execution with domain-specific compilers](https://mlsanity.com/paper/2102.13267/) (2021)
- [UNIT: Unifying Tensorized Instruction Compilation](https://mlsanity.com/paper/2101.08458/) (2021)

## Built on by (most cited first, 3 of 3)

- [ERASE: EaRly bAckpropagation SchEdule for Faster Training of Modern Recommendation Systems](https://mlsanity.com/paper/2608.18469/) (2026)
- [GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2](https://mlsanity.com/paper/2509.16248/) (2025)
- [Hybrid JIT-CUDA Graph Optimization for Low-Latency Large Language Model Inference](https://mlsanity.com/paper/2604.23467/) (2026)

## Abstract (arXiv)

Machine learning (ML) workloads launch hundreds to thousands of short-running GPU kernels per iteration. With GPU compute throughput growing rapidly, CPU-side launch latency of kernels is emerging as a bottleneck. CUDA Graphs promise to address this by replaying a set of kernels with a single dispatch of the graph, removing per-kernel launch costs. However, CUDA Graphs remain surprisingly difficult to deploy correctly and efficiently. We present PyGraph - a compiler framework to maximize the coverage and benefits of CUDA Graphs for ML workloads. It introduces three novel optimizations: it applies automatic code transformations to make ML applications amenable to CUDA Graphs; it eliminates the parameter copy overheads for kernels executing in CUDA Graphs, and it selectively deploys CUDA Graphs guided by a cost-benefit analysis. For 25 ML workloads from TorchBench, HuggingFace, and TIMM, PyGraph more than doubles the benefit from deploying CUDA Graph compared to the most popular and widely used ML compiler, PyTorch2. PyGraph is built atop PyTorch2's compilation framework and requires no programmer intervention.