arXiv:2503.19779 (cs) [2025年3月25日提交]

标题: PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch

作者: Abhishek Ghosh, Ajay Nayak, Ashish Panwar, Arkaprava Basu

查看由 Abhishek Ghosh 和其他 3 位作者撰写的名为 PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch 的论文 PDF 版本。 查看 PDF HTML (实验性)

摘要: CUDA Graphs 是 NVIDIA GPU 最近推出的一项硬件功能，旨在通过捕获和启动一系列 GPU 任务（kernels）作为一个 DAG 来减少 CPU 启动开销。然而，由于图的静态结构，部署 CUDA Graphs 目前面临着一些挑战。它还会因为数据复制而产生性能开销。事实上，我们展示了一个违反直觉的结果——在许多情况下，部署 CUDA Graphs 会损害性能。我们引入了 PyGraph，这是一种新颖的方法，可以自动利用 PyTorch2 中的 CUDA Graphs 的强大功能。受到三个关键观察结果的驱动，PyGraph 包含了三个新颖的优化：它实现了 CUDA Graphs 更广泛的部署，减少了 GPU kernel 参数复制开销，并基于成本效益分析有选择地部署 CUDA Graphs。PyGraph 与 PyTorch2 的编译工具链无缝集成，无需手动修改代码即可高效使用 CUDA Graphs。我们评估了 PyGraph 在各种机器学习基准测试中的表现，证明了相对于 PyTorch2 的显著性能改进。

主题： | 机器学习 (cs.LG) ---|--- 引用为： | arXiv:2503.19779 [cs.LG] (或 arXiv:2503.19779v1 [cs.LG] 用于此版本) https://doi.org/10.48550/arXiv.2503.19779 通过 DataCite 了解更多 arXiv 发布的 DOI

提交历史

来自：Abhishek Ghosh [查看电子邮件] [v1] 星期二, 25 Mar 2025 15:47:54 UTC (1,406 KB)

全文链接：

访问论文：

查看由 Abhishek Ghosh 和其他 3 位作者撰写的名为 PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch 的论文 PDF 版本。