数据科学

CUDA Python 1.0:稳定的 API、一个基础、完整的平台访问

多年来,需要 GPU 的 Python 开发者有两种现实选择:一是充分学习 NVIDIA CUDA C++ 以编写扩展程序,二是设置构建工具链,三是维护与 Python 的绑定 (大多数人从未这样做过) ;二是向上移动堆栈并让其他人的库 (即 PyTorch、CuPy 或 RAPIDS) 执行此操作。

第二个选择是 Python GPU 生态系统蓬勃发展的原因。但它也有局限性。当你不需要在上面的库中展示一些东西时,你就会回到第一选择。

由于每个库都以不同的方式访问 CUDA,因此让其中两个库在相同的数据上进行协作非常重要。如果 CuPy 分配了一个 GPU 显存块,cuDF 需要多少时间才能在同一流上处理该块而不进行复制?答案是通过交换协议和密切关注谁拥有什么。

在 CUDA 13.3 中,我们发布了 CUDA Python 1.0,这些库和工具可为 Python 提供完整的 CUDA 平台。Python 现已成为受支持的 CUDA 平台使用方式。

以下是所有内容:

  • cuda.core 1.0.0,Pythonic 访问 CUDA 运行时
  • cuda.compute 1.0.0,CCCL 的并行算法,可从 Python 调用
  • cuda.bindings 13.3.0,与 CUDA C API 的低级别 1:1 绑定,版本控制为 CUDA 工具包
  • cuda-pathfinder,用于查找环境中安装的 CUDA 组件
  • nvmath-python 1.0,Python 中的 NVIDIA 数学库,在自己的发布轨道上具有相同的稳定性承诺

CUDA Python 1.0 将命名为里程碑,而非您要在 pip 中输入的版本号;这些组件是独立版本控制的,因此上述不匹配的数字是故意的。

最重要的条目是 cuda.core。在这里,CUDA 的基本词汇 (设备、流、缓冲区) 变成了一组普通的 Python 对象,这一点远超便利:它为 Python 中的每个 GPU 库提供了一个共同的基础,以便在此基础上进行构建、协作和共享资源。这个理念是贯穿本文其余部分的线程。

CUDA 1.0:语义版本控制

CUDA Python 1.0 不是重写,也不是新产品。其中大多数库已经可用并改进了一段时间。1.0 带来的变化是一种承诺:语义版本控制

在实践中,这意味着:

  • 中断 API 更改仅在主要版本中发生
  • 次要版本添加了新功能
  • 补丁版本修复了错误
  • 任何计划移除的公共 API 都会在次要版本中被弃用,并且具有明确的替换路径

如果您因不确定 API 是否能在下次升级后继续使用而对构建库犹豫不决,那么这项保证就是标题。CUDA Python 将根据可预测的版本控制和弃用规则,在发布时持续跟踪新的 CUDA 功能。

一个基础而不是许多

要了解 1.0 版本的变化,记住之前的内容会有所帮助。

从 Python 访问 CUDA 过去意味着选择一个绑定层,其中有几个层。每个项目都由不同的项目进行维护,每个项目涵盖不同的 API 片段,每个项目都对流、设备或分配有自己的想法。如果您编写了应用程序,则继承了依赖项所使用的任何层。如果你编写库,你要么采用别人的库,要么构建自己的库,生态系统积累了更多的库。这就是改变。

现在,您可以通过一种由 NVIDIA 维护的官方方式从 Python 访问 CUDA。从 CUDA 13.3 开始,CUDA Python 和 C++ 等同于一等公民,NVIDIA 承诺在未来保持功能完全的奇偶校验。Python 是受支持的 CUDA 平台使用方式。

实际的好处是,库现在可以合成,而不仅仅是共存。Numba 内核和 cuda.compute 调用可以在同一流中的同一 GPU 缓冲区上运行,因为二者都没有引入私有 CUDA 层。对象跨越库边界,因为在下方,它们是相同的对象。对于共享问题,这比任何交换协议都要简短。

它还改变了使用高级平台功能的人。以前,像绿色上下文这样的功能会将 GPU 的流式多处理器进行分区,从而使延迟敏感型内核免受吞吐量内核的影响,因此每个感兴趣的库都需要将其独立绑定和公开。现在,它只需登陆一次 cuda.core,基于 cuda.core 构建的所有内容都可以访问它。

如果您构建面向 CUDA 的库,这将改变您的日常工作:您的工作将投入到使您的库与众不同的地方,而不是其他人已经编写的低级层。如果您编写应用程序,其优势将转移到一个级别,因为您的依赖项会在相同的流程上收。

思维模型:在一个基础上构建三层

CUDA Python 是涵盖 Python CUDA 生态系统的库集合:低级驱动和运行时绑定、并行算法、数学库、通信库和内核创作工具。下图 1 显示了它们的叠加情况。最清晰的图像读取方式是从下往上看。

运行时系统是其他所有系统所依赖的基础:设备管理、内存分配、流和同步、CUDA 计算图和 JIT 编译。

cuda-pathfinder 听起来很平淡,除非您记得 Python GPU 社区在调试进程实际加载的 CUDA 运行时方面花费了多少时间。

在此之上是 CUDA 库:Pythonic 可连接经 NVIDIA 调优的主机和设备库。cuda.compute 带来了 CCCL 的可主机调用并行算法;nvmath-python 现在也是 1.0,带来了包含主机 API、设备 API 和低级绑定的数学库;NCCL4Py 和 NVSHMEM4P 带来了通信库、NCCL 和 NVSHMEM。

它们都不会重新实现下方的 CUDA 层。它们在您直接使用的相同 cuda.core 缓冲区、设备和流中工作,例如,NVSHMEM 的对称内存会作为 cuda.core 缓冲区返回给您。共享基础不仅仅是对生态系统的指导;NVIDIA 自己的库也建立在此之上。

最重要的是核函数创作,以便您在需要自行编写 GPU 代码时使用。numba-cuda 是 Python 内核的 SIMT 语言,与之并行的是两种较新的特定领域语言:用于块模型的 CUDA 平铺语言 cutile-python 和用于 Tensor Core 的 CUTLASS 语言 cuteDSL

有两点需要澄清:

首先,您进入问题所需的级别,而无需全面学习。大量高效用户永远不会离开库层。图表中没有任何内容是课程。

其次,层并非单一版本产品。每个组件都会按自己的轨道移动,一些组件 (包括一些较新的内核创作语言) 仍处于试验阶段,尚未涵盖在 1.0 语义版本保证中,尽管其目的是让它们在稳定运行时处于相同的承诺之下。在将生产依赖关系固定为生产依赖关系之前,您需要了解这一点。

三种方式

在 CUDA Python 中,几乎每个人都会提出三个问题中的一个,并且每个问题都指向图表的不同层。在这里,它们的顺序是您承担的工作量,而不是它们在图表中的位置。

我只想优化算法:cuda.compute

最快的获胜方式通常是根本不编写核函数。它是一个已经存在的模型,并且已经被全职工作的人调整过。

cuda.compute 将 CUDA 核心计算库 (CCCL) 并行算法作为可主机调用的构建块引入 Python:排序、扫描、归约、转换、唯一、直方图、top-k 等。这些算法与支持高性能 C++ CUDA 代码的算法相同,而从 Python 来看,它们是对已有 GPU 数组的普通函数调用。事实证明,大部分数值运算都是这些模式的合成。由于 cuda.compute 为您实际运行的 GPU 编译它们,因此在新一代硬件上相同的调用不会改变。

版本 1.0 还使其更具表现力:您现在可以自定义算法使用普通 Python 函数 (包括 lambda) 的功能。

  • 在以下情况下选择此方案:您的问题分解为众所周知的并行模式,并且您希望使用最少的新代码获得结果。

我想用 Python 编写自己的内核:Numba

有时,您的逻辑与预打包算法不匹配。在这种情况下,您可以自己编写内核,并且仍然可以使用 Python 进行编写。

Numba 将 Python 子集编译为 GPU 内核:您可以使用装饰器标记函数,然后 Numba 从中生成 GPU 代码。您正在编写的是 CUDA SIMT 模型中的核函数,表示 GPU 一次在数千个线程上运行的单个线程的工作。进行这种思维转变是我们要学习的主要内容,而这一步远比使用 C++ 和构建系统小得多。

Numba CUDA MLIR 是一款兼容 Numba 的全新内核生成器,基于 MLIR 和现代 NVVM 工具链构建。它保留您已知的编程模型,并替换下方的编译器,从而提供更快的热 JIT 编译和更低的核函数启动延迟;对于大多数代码,迁移到它是单行导入更改。它比 1.0 组件更新,并且尚未包含在相同的语义版本控制承诺中。

  • 如果您的计算不适合预封装算法,并且您希望直接控制每个线程的功能,就可以实现这一点。

我需要驱动程序和运行时 API:cuda.corecuda.bindings

在基础部分,有两个软件包为您提供 CUDA 本身,而 cuda.core 达到 1.0 是此版本的核心。

cuda.core 是 CUDA 运行时的 Pythonic 接口,涵盖设备、流、程序、链接器、内存资源和图形,以及 CUDA C++ 的运行时编译,因此核函数无需单独的构建步骤即可从源代码运行到运行。重点在于 Pythonic:资源是真实的 Python 对象,故障会引发异常,而不是返回错误代码。由于它使用标准 CUDA 上下文,因此会与 Python GPU 生态系统的其余部分共享设备、流和内存,因此您自己的内核可以在 CuPy 数组或 PyTorch 张量上运行,而无需复制数据。

版本 1.0 将在之前发布周期中稳定运行的 API 整合到一个受支持的界面中,并添加了三种值得调用的功能:

  • 绿色上下文:如前所述,将 GPU 的 SM 划分为不相交组,以便在同一进程中避免延迟敏感型内核受到长时间运行的吞吐量内核的影响。
  • 进程检查点:快照显示正在运行的进程的完整 CUDA 状态,并在稍后恢复。
  • 进程间共享 (IPC) :在进程之间共享 GPU 显存,而无需通过主机进行复制。

在其下方,cuda.bindings 提供低级别绑定,可从驱动和运行时通过其周围的编译器、链接器和系统库,以 1:1 的比例全面覆盖 CUDA 主机 API。它采用 CUDA 工具包进行版本控制。cuda.core 针对 Python 人体工程学进行了优化,而 cuda.bindings 则针对完整性进行了优化:如果存在于 C API 中,则可以从 Python 访问。

  • 在以下情况下选择此方案:当您要构建 GPU 库或将 CUDA 集成到现有工具包中使用 cuda.core 提高 Python 效率,并在需要全面访问 C API 时使用 cuda.bindings

生态系统已经在融合

谁已经在共享基础上进行构建,是共享基础有效运作的最好证明。如前文所述,NVIDIA 通信库和数学库已经全部支持 cuda.core 对象。更广泛的生态系统也在融合。在导入模块时,CuPy 的构建更简单,占用空间更小,速度更快。PyTorch 的 CUDA wheel 现在依赖于 cuda.bindings

移动到共享层上的每个库都会从依赖关系图中再提取一个私有绑定层。这减少了版本冲突,减少了神秘的互操作错误,并且减少了两个库在各自的 CUDA 环境中存在不同意见的地方。

而且,由于每个组件都在自己的轨道上遵循语义版本控制,因此依赖其中一个组件并不意味着继承其余组件的数据流。

开始使用

只需一条命令,即可获取 CUDA Python 堆栈:

pip install cuda-python cuda-cccl numba-cuda-mlir[cu13]

本文将介绍上述 CUDA Python 组件,以及基于 MLIR 的 Numba 后端。使用 pip install nvmath-python[cu13] 单独安装 nvmath-python。唯一的系统要求是更新的 NVIDIA 驱动程序;通常不需要单独安装 CUDA 工具包。

如果您决定从何处着手:

  • 如果您的工作是数据科学,您可能根本不需要达到这种低水平。RAPIDS 库已经覆盖了这一领域:cuDF 可加速 pandas、Polars 和 Apache Spark,而 nx-cugraph 支持 NetworkX
  • 如需优化算法,请从 cuda.compute 开始
  • 要使用 Python 编写自己的核函数,请从 Numba 或 Numba CUDA MLIR 开始
  • 要访问低级驱动和运行时 API,请从 cuda.corecuda.bindings 开始

之后,CUDA Python 文档NVIDIA/cuda-python 存储库提供安装指南、API 参考和示例,而NVIDIA 加速计算中心则会收集更广泛的 GPU 计算学习材料。

1.0 里程碑中最精彩的部分在于,这一决定不再是高风险。选择与您面前的问题相匹配的等级,并确定其下方的地面是稳定的。

致谢

CUDA Python 1.0 反映了 CUDA Python 产品和工程团队多年来的工作成果,他们设计和构建了此处所述的库。还要感谢 NVIDIA 的评审人员,他们的反馈改进了版本和本博文,还要感谢提交问题、测试预发行版并帮助塑造我们现在承诺支持的 API 的开源贡献者。

标签