PyGraph：PyTorch 中 CUDA Graphs 的鲁棒编译器支持
机器学习
2025-12-24 v3
摘要
机器学习（ML）工作负载在每次迭代中启动数百至数千个短时运行的 GPU kernel。随着 GPU 计算吞吐量的快速增长，CPU 侧的 kernel 启动延迟正成为瓶颈。CUDA Graphs 旨在通过以单次图分发重放一组 kernel 来解决此问题，从而消除每 kernel 的启动开销。然而，正确且高效地部署 CUDA Graphs 仍然出人意料地困难。我们提出了 PyGraph——一个旨在最大化 ML 工作负载中 CUDA Graphs 覆盖范围与收益的编译器框架。它引入了三项新颖的优化：应用自动代码转换使 ML 应用程序适配 CUDA Graphs；消除在 CUDA Graphs 中执行的 kernel 的参数拷贝开销；并在成本效益分析的指导下选择性部署 CUDA Graphs。对于来自 TorchBench、HuggingFace 和 TIMM 的 25 个 ML 工作负载，与最流行且广泛使用的 ML 编译器 PyTorch2 相比，PyGraph 将部署 CUDA Graph 的收益提高了一倍以上。PyGraph 构建于 PyTorch2 的编译框架之上，无需程序员干预。
引用
@article{arxiv.2503.19779,
title = {PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch},
author = {Abhishek Ghosh and Ajay Nayak and Ashish Panwar and Arkaprava Basu},
journal= {arXiv preprint arXiv:2503.19779},
year = {2025}
}