The CUDA Python Developer's Toolbox

Katrina Riehl, PhD
Principal Technical Product Manager - CUDA Python
GTC 2025 - March 18, 2025

目录

1. CUDA与Python：光速生产力的故事

Page 2
Page 2

本次演讲是关于CUDA和Python系列讲座的一部分。该系列共包含四个步骤，旨在引导开发者逐步深入了解并掌握在Python中使用CUDA进行加速计算的技术：

  • 第0步：聆听我们的Python故事 (Hear Our Python Story)

    • 加速Python：社区与生态系统 (Accelerated Python: The Community and Ecosystem)
    • 时间：周二, 3月18日, 3:00 PM - 3:40 PM PDT

  • 第1步：学习CUDA Python工具 (Learn the CUDA Python Tools)

    • CUDA Python开发者工具箱 (The CUDA Python Developer's Toolbox)
    • 时间：周三, 3月19日, 10:00 AM - 10:40 AM PDT

  • 第2步：深入了解内核 (Dive Deep into Kernels)

    • 在Python中编写CUDA内核的1001种方法 (1,001 Ways to Write CUDA Kernels in Python)
    • 时间：周三, 3月19日, 11:00 AM - 11:40 AM PDT

  • 第3步：精通Tensor Core (Master the Tensor Core)

    • 使用CUTLASS 4.0在Python中实现Tensor Core编程 (Enable Tensor Core Programming in Python With CUTLASS 4.0)
    • 时间：周五, 3月21日, 11:00 AM - 11:40 AM PDT

2. CUDA加速Python：贯穿整个技术栈的一流Python支持

"一流"（First-Class）意味着CUDA现在将Python融入到平台的每一个API和每一个特性中。

Page 4
Page 4

CUDA Python生态系统提供了一个分层的技术栈，涵盖了从高层框架到底层工具的全面支持。这些层次之间具有互操作性，允许开发者混合和匹配不同层次的库来满足其需求。

技术栈各层级及其对应的Python库示例：

  • 框架与领域特定语言 (Frameworks & DSLs): Warp, Omniverse, PyTorch
  • 软件开发工具包 (SDKs): RAPIDS, CUDA-Q, Holoscan
  • 领域特定库 (Domain-Specific Libraries): cuPyNumeric, CuPy
  • 加速库 (Accelerated Libraries): cuda.parallel, nvmath-python
  • 通信库 (Communication Libraries): mpi4py, nvshmem4py
  • 设备库 (Device Libraries): cuda.cooperative, nvmath-python
  • 内核编写 (Kernel Authoring): numba-cuda, cuTe, CUTLASS DSL
  • 编译器栈 (Compiler Stack): numba-ir, nvjitlink
  • 主机运行时与工具 (Host Runtimes & Tools): cuda.bindings, cuda.core, Nsight Systems

CUDA Python 框架

PyTorch 是一个开源的机器学习（ML）框架。以下代码展示了如何使用PyTorch实现SAXPY（y = a*x + y）操作。

import torch
def saxpy(a:float, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Perform the SAXPY operation: y = a * x + y
    """
    
    return a * x + y

# Example usage:
a = 2.0
x = torch.tensor([1.0, 2.0, 3.0])
y = torch.tensor([4.0, 5.0, 6.0])

result = saxpy(a, x, y)
print(result) # Output: tensor([6., 9., 12.])

CUDA Python SDKs

CuDF 是一个类似于GPU上的Pandas的库，它是RAPIDS生态系统的一部分。

Page 6
Page 6

以下代码展示了如何使用cuDF创建一个DataFrame并进行操作。

df = cudf.DataFrame(
    {'angles': [0, 3, 4], 'degrees': [360, 180, 360]},
    index=['circle', 'triangle', 'rectangle']
)

>>> df.multiply(1)

#          angles degrees
# circle        0     360
# triangle      3     180
# rectangle     4     360

CUDA Python 领域特定库

CuPy 是一个类似于GPU上的NumPy和SciPy的库。

Page 7
Page 7

以下代码展示了如何使用CuPy在GPU上创建数组、执行计算，并将结果移回主机（CPU）。

x_gpu = cp.array([1, 2, 3]) # create an array in the current device
y_gpu = cp.array([1, 2, 3])

x_gpu = x_gpu + y_gpu

x_cpu = cp.asnumpy(x_gpu) # move the array to the host

CUDA生态系统的不同层面

CUDA Python库在Python开发者体验中的定位可以在生产力（productivity）与速度（speed）的权衡曲线上体现。通常，更高层次的抽象（如框架）提供更高的生产力但可能对性能的控制较少，而更低层次的库则提供更高的速度和控制力，但开发复杂度更高。CUDA Python旨在提升整个曲线，让开发者在获得高性能的同时也能保持高生产力。

Page 9
Page 9

3. 深度学习用例：深入研究激活函数

Page 10
Page 10
  • 像PyTorch这样的深度学习框架利用GPU内核来实现激活函数，这是神经网络的一个关键组成部分，能够实现更快的训练和推理。

  • 这些内核利用GPU的并行处理能力来执行以下操作：

    • 逐元素计算（element-wise calculations）
    • 矩阵乘法（matrix multiplication）
    • 卷积（convolutions）

  • 本次演讲将使用修正线性单元 (Rectified Linear Unit, ReLU) 激活函数（包括带偏置和不带偏置的情况）作为示例。

使用CuPy实现

我们将首先关注技术栈中的CuPy库。

Page 11
Page 11

CuPy是一个与NumPy/SciPy兼容的数组库，用于通过Python进行GPU加速计算。

以下示例展示了如何使用CuPy实现一个带有偏置的ReLU函数：d = ReLU(a @ b + bias)。

Page 12
Page 12

  • 实现步骤：

    1. 矩阵乘法 (Matrix Multiplication)
    2. 加上偏置 (Add Bias)
    3. 对小于0的值进行掩码操作 (Mask for values < 0)

  • 代码分析：

    • CuPy的有状态API（Stateful API）将计划（planning）与执行（execution）分开。
    • mm = Matmul(weights, x) 和 mm.plan() 阶段用于规划矩阵乘法操作。
    • 在forward函数中，y = mm.execute() 执行计算，然后加上偏置，最后通过 y[y < 0] = 0 实现ReLU功能。

CuPy 亮点

Page 13
Page 13
  • 开源: https://github.com/cupy/cupy
  • 可在 PyPI 和 conda-forge 上获取。
  • 支持 CUDA 11 和 12。
  • 非NVIDIA CUDA特定: 这是一个独立的项目。

  • 目标:

    • CuPy项目的目标是为Python用户提供GPU加速能力，而无需深入了解底层GPU技术。
    • CuPy团队专注于提供：
      • 完整的NumPy和SciPy API覆盖，以作为直接替代品，并提供高级CUDA功能以最大化性能。
      • 一个成熟且高质量的库，作为所有需要加速的项目的基本包，从实验室环境到大规模集群。

  • 特性:

    • 与CPU上的NumPy、SciPy具有互操作性。
    • 与nvmath-python、RAPIDS和PyTorch具有互操作性。

使用 nvmath-python 实现

接下来，我们将关注技术栈中的 nvmath-python 库。

Page 14
Page 14

CUDA-X 数学库 (C++层面)

nvmath-python 是对底层 CUDA-X 数学库的封装。这些由NVIDIA提供的GPU加速数学库是计算密集型应用的基础，应用领域涵盖分子动力学、计算流体动力学、计算化学、医学成像以及地震勘探等。

Page 15
Page 15

主要的CUDA数学库包括：

  • cuBLAS: GPU加速的基础线性代数 (BLAS) 库。
  • cuFFT: GPU加速的快速傅里叶变换 (FFT) 库。
  • cuRAND: GPU加速的随机数生成库。
  • cuSOLVER: GPU加速的稠密和稀疏直接求解器。
  • cuSPARSE: GPU加速的稀疏矩阵库。
  • cuTENSOR: GPU加速的张量线性代数库。
  • cuDSS: GPU加速的直接稀疏求解器库。
  • CUDA Math API: GPU加速的数学函数API。
  • AmgX: 用于模拟和隐式非结构化方法的GPU加速线性求解器。

nvmath-python: RELU 示例

nvmath-python 是一个用于在 Python 中利用 NVIDIA GPU 加速数学计算的库。以下通过一个 d = RELU(a @ b + bias) 的例子，展示其核心特性。

nvmath-python 及其特性展示，包含代码、日志和性能图表 (Page 16)
nvmath-python 及其特性展示，包含代码、日志和性能图表 (Page 16)

核心特性：

  • 与其他库/框架的互操作性：可以与 PyTorch 等主流框架无缝集成，直接使用 torch.rand 创建张量。
  • 支持多种 epilogues：在主计算（如矩阵乘法）之后，可以附加多种操作（如加偏置和 RELU），并将所有操作融合（Fuses all operations）到一个单一的 CUDA 核中，以提升性能。
  • 支持日志记录以获取深入洞察：通过开启日志，可以详细了解内部操作，如操作类型、数据类型、内存限制、批处理大小、矩阵维度等。
  • 提供对象 API 以实现对行为的最佳控制：用户可以通过面向对象的接口精细控制计算过程。
  • 允许内核自动调优以获得最佳性能：库可以自动测试不同的算法方案，并选择性能最佳的配置。如日志所示，它会评估多个方案并根据启发式规则选择最优的四个进行测试，最终确定性能最好的方案。

性能对比：

图表显示了 d = RELU(a @ b + bias) 在不同实现下的性能（以 TFLOPS 计）。
* 使用三个独立的内核（3 kernels）性能约为 76.5%。
* 使用一个默认启发式选择的融合内核（1 fused kernel, default heuristic）性能约为 84.2%。
* 使用自动调优（Autotuned）找到的最佳融合内核性能最高，达到 89.9%。

nvmath-python 亮点

nvmath-python 亮点及二维码 (Page 17)
nvmath-python 亮点及二维码 (Page 17)
  • 开源：代码托管在 GitHub: https://github.com/nvidia/nvmath-python
  • 包管理：可通过 PyPI 和 conda-forge 安装。

  • 环境支持：

    • 支持 CUDA 11 和 12。
    • 支持 Linux (x86-64, aarch64) 和 Windows (x64)。

  • 目标用户：

    • 寻求生产力、与其他库和框架的互操作性以及高性能的研究人员。
    • 寻求开箱即用的性能和通过 Python 获得更好可维护性的库/框架开发者。
    • 寻求最高性能而无需切换到 CUDA 的内核开发者。

  • Beta 版本特性：

    • 提供对 CUDA 数学库的低级绑定。
    • 提供 Pythonic 的高级 API（主机端和设备端）。
    • 目前该版本已扩展到矩阵乘法（matmul）和快速傅里叶变换（FFTs）。
    • 支持在 Numba 内核中进行设备函数调用。
    • 与 NumPy, CuPy 和 PyTorch 张量具有互操作性。

CUDA 加速的 Python：贯穿整个技术栈的一流 Python 支持

“一流”意味着 CUDA 现在将 Python 纳入其平台的每个 API 和每个功能中。下图展示了 CUDA Python 生态系统在不同层级的分布，其中 numba-cuda 位于内核编写（Kernel Authoring）层。

CUDA Python 技术栈概览 (Page 18)
CUDA Python 技术栈概览 (Page 18)

Numba CUDA

Numba 是一个即时（Just-In-Time, JIT）编译器，用于加速 Python 中的数值函数在 CPU 和 GPU 上的运行。它允许开发者直接在 Python 中编写 CUDA 内核。

Numba CUDA 代码示例 (Page 19)
Numba CUDA 代码示例 (Page 19)

代码示例：

  • 左侧代码：一个简单的一维向量乘法内核。它使用 cuda.grid(1) 获取线程索引，并直接操作 NumPy 数组。Numba 负责处理线程和块的索引。
  • 右侧代码：一个扩展到二维的矩阵乘法内核。它使用 cuda.grid(2) 获取二维的行和列索引，实现了标准的矩阵乘法逻辑。

Numba CUDA: GPU 上的 ReLU 内核

以下是使用 Numba CUDA 实现 ReLU 激活函数的内核示例。

Numba CUDA 实现的 ReLU 内核代码 (Page 20)
Numba CUDA 实现的 ReLU 内核代码 (Page 20)
  • 输入/输出：内核接受 NumPy 或 CuPy 数组作为输入和输出。
  • 一维实现：通过 cuda.grid(1) 获取全局线程 ID (idx)，并对一维数组的每个元素应用 max(0.0, element) 操作。
  • 二维实现：通过 cuda.grid(2) 获取行 (row) 和列 (col) 索引，对二维数组的每个元素进行操作，逻辑与一维版本类似，但扩展到了二维空间。

Numba 概览

Numba 概览及二维码 (Page 21)
Numba 概览及二维码 (Page 21)
  • 开源：代码托管在 GitHub: https://github.com/NVIDIA/numba-cuda
  • 包管理：可通过 PyPI 和 conda-forge 安装。

  • 环境支持：

    • 支持 CUDA 11 和 12。
    • 支持 Linux (x86-64, aarch64) 和 Windows (x64)。

  • 核心功能：Numba 通过将受限的 Python 代码子集直接编译为 CUDA 内核和设备函数来支持 CUDA GPU 编程，遵循 CUDA 执行模型。

  • CUDA C 映射：Numba 暴露的大多数 CUDA 编程设施都直接映射到 NVIDIA 提供的 CUDA C 语言。
  • 目标用户：
    • Python 科学计算开发者、研究人员和数据科学家。
    • 寻求最高性能而无需切换到 CUDA 的内核开发者。

CUDA 加速的 Python 技术栈及其互操作性

CUDA Python 技术栈的各个层级之间可以混合搭配，具有良好的互操作性。例如，开发者可以在 PyTorch（框架层）中使用 CuPy（领域特定库），调用 nvmath-python（加速库），甚至可以在其中嵌入使用 numba-cuda（内核编写）编写的自定义内核。

CUDA Python 技术栈的互操作性 (Page 22)
CUDA Python 技术栈的互操作性 (Page 22)

4. CUDA开发者工具

为了在 GPU 上实现光速般的性能，可以使用 NVIDIA 提供的开发者工具。

CUDA 开发者工具及性能优化流程 (Page 23)
CUDA 开发者工具及性能优化流程 (Page 23)

性能优化考虑因素：

  • 在 CUDA 中实现最佳性能：

    • 本地化内存访问。
    • 最大化活动线程数。
    • 最小化条件分支。

  • 克服 CPU 和 GPU 间的 PCIe 瓶颈：

    • 最小化数据传输量。
    • 组织数据以补充硬件架构。
    • 利用异步传输。

开发者工具：

  • 使用 Nsight Systems 和 Nsight Compute 检测瓶颈和性能缺陷：

    • Nsight Systems：用于端到端的综合性能分析。
    • Nsight Compute：用于 CUDA 内核的详细性能分析。
    • 这是一个循环过程：通过 Nsight Systems 进行宏观分析，然后用 Nsight Compute 进行深度分析，优化后再重新检查。

  • 使用 Compute Sanitizer 保证正确性：

    • Memcheck (默认)：检测越界和内存访问错误。
    • Racecheck：检测竞争条件。
    • Initcheck：检测未初始化的内存访问。
    • Synccheck：检测导致死锁或竞争条件的无效同步模式。

CUDA Python 开发者工具特性

这些工具超越了传统的 pdb 和 print 调试方式，为 Python 开发者提供了强大的功能。

  • Nsight Systems

    • 全局解释器锁（GIL）追踪。
    • C++ 和 Python 帧的组合回溯。
    • 使用 NVIDIA Tools Extension (NVTX) API 注释代码。

  • Nsight Compute

    • Numba 内核调试。

  • Nsight Systems 和 Nsight Compute 的共同特性

    • JupyterLab 集成。
    • 与 CUPTI 绑定的集成。

开发者工具生态系统

NVIDIA 提供了一个完整的工具生态系统，以帮助开发者在 GPU 上实现极致性能。

开发者工具生态系统概览 (Page 25)
开发者工具生态系统概览 (Page 25)
  • 调试器 (Debuggers)：cuda-gdb, Nsight Visual Studio Edition, Visual Studio Code。
  • 性能分析器 (Profilers)：Nsight Systems, Nsight Compute, CUPTI, NVIDIA Tools eXtension (NVTX)。
  • 正确性检查器 (Correctness Checker)：Compute Sanitizer。
  • IDE 集成 (IDE Integrations)：Nsight Visual Studio Edition, Nsight Eclipse Edition。

5. 社区资源与相关会议

加速计算中心 (Accelerated Computing Hub)

NVIDIA 在 GitHub 上提供了一个加速计算中心，汇集了与通用和加速 GPU 编程相关的教育资源。

访问地址: https://github.com/NVIDIA/accelerated-computing-hub

NVIDIA 加速计算中心 GitHub 页面 (Page 26)
NVIDIA 加速计算中心 GitHub 页面 (Page 26)

CUDA 开发者会议

以下是与 CUDA 相关的更多开发者会议列表，涵盖了通用 CUDA、CUDA Python、CUDA C++、开发者工具、多 GPU 编程和性能优化等多个主题。

CUDA 开发者会议列表 (Page 28)
CUDA 开发者会议列表 (Page 28)

  • 通用 CUDA

    • S72571 - What's CUDA All About Anyways?
    • S72897 - How To Write A CUDA Program: The Parallel Programming Edition

  • CUDA Python

    • S72450 - Accelerated Python: Tour of the Community and Ecosystem
    • S72448 - The CUDA Python Developer's Toolbox
    • S72449 - 1001 Ways to Write CUDA Kernels in Python
    • S74639 - Enable Tensor Core Programming in Python With CuTe

  • CUDA C++

    • S72574 - Building CUDA Software at the Speed-of-Light
    • S72572 - The CUDA C++ Developer's Toolbox
    • S72575 - How You Should Write a CUDA C++ Kernel

  • 开发者工具

    • S72527 - It's Easier than You Think - Debugging and Optimizing CUDA with Intelligent Developer Tools

  • 联系专家

    • CWE72433 - CUDA Developer Best Practices
    • CWE73310 - Using NVIDIA CUDA Compiler Tool Chain for Productive GPGPU Programming
    • CWE72393 - What's in Your Developer Toolbox? CUDA and Graphics Profiling, Optimization, and Debugging Tools

  • 多 GPU 编程

    • S72576 - Getting Started with Multi-GPU Scaling: Distributed Libraries
    • S72579 - Going Deeper with Multi-GPU Scaling: Task-based Runtimes
    • S72578 - Advanced Multi-GPU Scaling: Communication Libraries

  • 性能优化

    • S72683 - CUDA Techniques to Maximize Memory Bandwidth and Hide Latency
    • S72685 - CUDA Techniques to Maximize Compute and Instruction Throughput
    • S72686 - CUDA Techniques to Maximize Concurrency and System Utilization
    • S72687 - Get the Most Performance from Grace Hopper

更多信息请访问：nvidia.com/gtc/sessions/cuda-developer