PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch
Not on the map. Why?
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.
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
23 of 90 references.
References · 23
- 2021Learning Transferable Visual Models From Natural Language Supervisioncited by 33,744
- 2022Hierarchical Text-Conditional Image Generation with CLIP Latentscited by 6,585
- 2019XLNet: Generalized Autoregressive Pretraining for Language Understandingcited by 2,811
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
- 2017Mixed Precision Trainingcited by 1,156
- 2024Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Servecited by 526
- 2014Caffe: Convolutional Architecture for Fast Feature Embeddingcited by 288
- 2024InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Managementcited by 245
- 2020Ansor: Generating High-Performance Tensor Programs for Deep Learningcited by 201
- 2024Parrot: Efficient Serving of LLM-based Applications with Semantic Variablecited by 99
- 2019Analysis of Large-Scale Multi-Tenant GPU Clusters for DNN Training Workloadscited by 79
- 2018Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Codecited by 53
11 more
- 2022SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learningcited by 47
- 2024POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inferencecited by 47
- 2024Mirage: A Multi-Level Superoptimizer for Tensor Programscited by 41
- 2022TensorIR: An Abstraction for Automatic Tensorized Program Optimizationcited by 38
- 2022Tensor Program Optimization with Probabilistic Programscited by 30
- 2022Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programscited by 28
- 2020Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inferencecited by 19
- 2021Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performancecited by 18
- 2021The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Paddingcited by 17
- 2021LazyTensor: combining eager execution with domain-specific compilerscited by 11
- 2021UNIT: Unifying Tensorized Instruction Compilationcited by 10
Built on by
Citing papers per year: 2025: 1; 2026: 2.
3 papers.