Python 在符合人体工程学的机器学习领域占据主导地位，但编写真正高效的 GPU 代码历来需要使用 C++ 编写自定义内核，并维护与 Python 的绑定。对大多数 Python 开发者和研究人员而言，这构成了显著的入门障碍。
PyTorch 等框架通过在 CUDA C++ 中实现内核（手写或借助 NVIDIA CUDA 核心计算库 等库）来解决这一问题。手写内核十分耗时，且需要深厚的底层架构专业知识。使用 CUB（CCCL 中的 C++ 库）通常是更优的选择，因为其提供的基元在各个架构上均经过高度优化，并经过严格测试。然而传统上，将 CUB 与 Python 结合使用需要构建和维护绑定，并对 C++ 模板进行固定类型和运算符的预实例化，这在一定程度上限制了 Python 端的灵活性。
NVIDIA cuda.compute 库通过为设备范围的 CUB 基元提供高级的 Pythonic API 来克服这些限制。
借助 cuda.compute，NVIDIA CCCL 团队在 GPU MODE 排行榜上位居前列，这是一个由拥有超过 20,000 名成员的在线社区主办的内核竞赛，旨在促进 GPU 编程的学习与优化。GPU 模式通过举办内核竞赛，为各类计算任务（从简单的向量加法到更复杂的块矩阵乘法）寻找高效实现方案。
NVIDIA CCCL 团队专注于通过高级抽象实现跨 GPU 架构的并行基元的“光速” (SOL) 实现。该团队在已验证的 GPU 架构（NVIDIA B200、NVIDIA H100、NVIDIA A100 和 NVIDIA L4）上均展现出卓越的整体性能。
在这篇博客中，我们将分享更多细节，说明如何在排行榜上取得如此优异的排名。
CUDA Python：GPU 性能满足生产力需求
CUB 为常见并行操作（包括 GPU MODE 竞赛中采用的并行操作）提供高度优化的 CUDA 内核。这些内核经过架构调优，被广泛认为是接近极致性能的实现。
cuda.compute 库支持在 Python 中直接定义自定义类型和运算符。本质上，它能够即时（JIT）编译专用内核，并应用链接时优化，从而实现接近 SOL 的性能水平，与 CUDA C++ 相当。您可继续使用 Python，同时兼具模板的灵活性和经过优化的 CUDA 内核所具备的高性能。
借助 cuda.compute，您将获得：
- 在 Python 中构建快速、可组合的 CUDA 工作流：直接在 Python 中开发高效且模块化的 CUDA 应用。
- 支持自定义数据类型与运算符：无需依赖 C++ 绑定即可灵活使用自定义类型和操作。
- 实现高性能优化：基于经过验证的 CUB 基元，充分发挥硬件架构潜力。
- 加速迭代开发：借助即时编译（JIT）技术，在保持 CUDA C++ 级性能的同时提升开发效率。JIT 编译在不牺牲性能的前提下，为开发者提供更高的灵活性与更快的迭代速度，显著缩短开发周期。
排行榜结果
我们使用 cuda.compute 提交了 GPU MODE 基准测试中 PrefixSum、VectorAdd、Histogram、Sort、和 Grayscale（查找用户名 Nader）。
对于排序等算法，CUB 的实现速度比次优提交快 2 到 4 倍。这就是 CCCL 的实际承诺：SOL+ 类算法在处理标准基元时的性能优于自定义内核，否则您需要花费几个月的时间来构建。
如果我们没有位居首位，差距通常源于缺乏针对该特定 GPU 的调优策略。在某些情况下，我们的实现是一种更为通用的解决方案，而排名靠前的提交则专门针对特定的问题规模进行了优化。
在其他情况下，排名第一的提交作品已在幕后使用 CUB 或 cuda.compute。这表明，这些库已经体现了众多标准 GPU 算法的性能极限，其性能特征如今已被充分理解，并被领先的提交者广泛依赖。
这与获胜无关
排行榜上的结果只是附带产物；真正的目标是与社区共同学习、透明地开展基准测试，并展现 Python 在高性能 GPU 任务中的强大能力。
我们的目标并非阻止手写 CUDA 内核。对于自定义内核（如新算法、紧密融合或专门的内存访问模式），存在大量合理的应用场景，但标准基元（排序、扫描、归约、直方图等）除外。对于这些标准基元，首选应是经过验证的高性能实现。借助 cuda.compute，这些经过调优的 CUB 基元现在可直接从原生 Python 访问，使您能够构建高质量、可投入生产的 GPU 加速 Python 库。
对于任何正在构建下一个 CuPy、RAPIDS 组件或自定义 Python GPU 加速库的人来说，这无疑是个好消息：在保持纯 Python 环境的同时，实现更快的迭代速度、更少的胶水层以及接近生产级的性能。
cuda.compute 在实践中的效果如何
任何人在学习 GPU 编程时，通常首先要编写的示例之一就是向量加法。使用 cuda.compute，我们可以通过调用设备范围的基元，用纯 Python 来解决这个问题。
import cuda.compute
from cuda.compute import OpKind
# Build-time tensors (used to specialize the callable)
build_A = torch.empty(2, 2, dtype=torch.float16, device="cuda")
build_B = torch.empty(2, 2, dtype=torch.float16, device="cuda")
build_out = torch.empty(2, 2, dtype=torch.float16, device="cuda")
# JIT compiling the transform kernel
transform = cuda.compute.make_binary_transform(build_A, build_B, build_out, OpKind.PLUS)
# Defining custom_kernel is required to submit to the GPU MODE competition
def custom_kernel(data):
# Invoking our transform operation on some input data
A, B, out = data
transform(A, B, out, A.numel())
return out
您可以在 GPU 模式排行榜 上找到更多 cuda.compute 示例。这种模式具有一致性：通过调用设备范围内的构建块（由 CCCL 针对每一代 GPU 自动优化），实现性能卓越的简洁代码。
其他针对 VectorAdd 类别的优秀提交需要引入 C++ 和内联 PTX，因此代码高度依赖于架构。
立即试用 cuda.compute
如果您要构建 Python GPU 软件、自定义工作流、库组件或对性能要求较高的代码，可以通过 cuda.compute 直接在 Python 中使用 CCCL CUB 基元，借助专为架构感知和极致性能优化而设计的构建块来提升效率。
要试用 cuda.compute，您可以通过 pip 或 conda 进行安装：
pip install cuda-cccl[cu13] (or [cu12])
conda install -c conda-forge cccl-python cuda-version=12 (or 13)
我们正与社区共同构建此框架，您的反馈和基准测试将影响我们的路线图，欢迎通过 Github 或 GPU MODE Discord 与我们联系。
