optimize-for-gpu

使用 NVIDIA 优化 Python 的 GPU 性能

应将 GPU 加速视为以证据为依据的优化，而不是自动重写。保留用户的数值和算法约定，使用有代表性的数据进行测量，并且只有在同步的端到端基准测试显示出实用的性能提升时，才保留 GPU 版本。

适用场景

• 用户希望加快数值或科学计算类 Python 代码
• 用户正在处理大型数组、矩阵或数据框
• 用户提到 CUDA、GPU、NVIDIA 或并行计算
• 用户的 NumPy、pandas、SciPy、scikit-learn、NetworkX 或 scipy.sparse.linalg 代码需要处理大型数据集
• 用户需要底层 GPU 基础操作（稀疏特征值求解器、设备内存管理、多 GPU 通信）
• 用户正在进行机器学习（训练、推理、超参数调优、预处理）
• 用户正在进行图分析（中心性、社区检测、最短路径、PageRank 等）
• 用户正在进行向量搜索、最近邻搜索、相似度搜索，或构建 RAG 流水线
• 用户有可通过 GPU 加速的 Faiss、Annoy、ScaNN 或 sklearn NearestNeighbors 代码
• 用户希望在大型数据集上使用 GPU 加速的交互式仪表板、交叉筛选或探索性数据分析
• 用户正在使用 GeoPandas 或 shapely 进行地理空间分析（点在多边形内判定、空间连接、轨迹分析、距离计算）
• 用户正在使用 scikit-image 或 OpenCV 进行图像处理、计算机视觉或医学影像处理（滤波、分割、形态学、特征检测）
• 用户正在处理全切片图像（WSI）、数字病理、显微成像或遥感影像
• 用户正在将大型二进制数据文件加载到 GPU 内存（numpy.fromfile → cupy，或 Python open() → GPU 数组）
• 用户需要将 S3、HTTP 或 WebHDFS 中的文件直接读入 GPU 内存
• 用户提到 GPUDirect Storage（GDS），或希望在文件 IO 时绕过 CPU 内存暂存
• 用户正在进行物理模拟（粒子、布料、流体、刚体）或可微分模拟
• 用户需要在 GPU 上执行网格操作（光线投射、最近点查询、有符号距离场）或几何处理
• 用户正在使用变换和四元数进行机器人学（运动学、动力学、控制）
• 用户有可通过 JIT 编译为 GPU 内核的 Python 模拟循环
• 用户提到 NVIDIA Warp，或希望将可微分 GPU 模拟集成到 PyTorch/JAX
• 用户正在进行模拟、信号处理、金融建模、生物信息学、物理学或任何计算密集型工作
• 用户希望优化现有代码，且 GPU 加速是合适的方案

选择最合适的最小层

优先使用维护良好的库实现，而非自定义内核：

现有工作负载	首选路径	用途
NumPy / SciPy	CuPy	数组、稀疏矩阵、线性代数、FFT、信号处理
pandas	cudf.pandas，然后 cuDF	优先使用加速器模式；需要更多控制时使用原生 API
scikit-learn	cuml.accel，然后 cuML	优先使用加速器模式；按需使用原生估计器
NetworkX	nx-cugraph，然后 cuGraph	优先使用后端分发；大规模场景使用原生图 API
scikit-image	cuCIM	GPU 图像处理和全切片成像
Faiss / Annoy / k-NN	cuVS	精确和近似向量搜索
原始或远程文件 I/O	KvikIO	GPU 缓冲区和 GPUDirect Storage
自定义数组内核	新项目使用 Numba-CUDA-MLIR；现有代码使用 Numba-CUDA	显式 SIMT 内核和共享内存
空间或可微分内核	Warp	几何、仿真内核、机器人、自动微分
高层物理仿真	Newton	接替已移除的 warp.sim 模块的维护良好的引擎
底层 RAPIDS 原语	RAFT（pylibraft）	稀疏特征求解器、多 GPU 构建模块

不要仅仅为了使用这些库而将代码从 PyTorch、JAX、TensorFlow 或其他 GPU 原生框架中迁出。 首先消除 CPU 往返，并使用该框架的编译器、性能分析器、混合精度和批处理功能。

将以下项目仅视为遗留项目：

项目	状态	指导
cuxfilter	最终版本 26.06	仅维护现有仪表板。对于新项目，将 cuDF 与 HoloViews/hvPlot/Datashader 结合，并使用 Panel、Dash、Streamlit 或 Bokeh 提供服务。
cuSpatial	已于 25.04 归档	仅在隔离的遗留环境中使用。对于新项目，将几何数据保留在 GeoPandas/Shapely 中，并使用 cuDF 加速兼容的表格阶段。

每个库的详细指导，包括每个库何时是错误选择以及如何组合使用它们，请参见 references/decision_framework.md。 安装命令和 CUDA 版本选择请参见 references/installation.md。每个库的前后转换示例请参见 references/code_transformation_patterns.md。

优化工作流

1. 定义契约和基线

• 捕获一个具有代表性的输入、预期输出和可接受的数值容差。
• 测量当前的端到端路径，包括输入、传输、计算和输出。
• 在更改代码之前进行性能分析。对 CPU 代码使用 CPU 性能分析器，并确定真正的瓶颈 是计算、内存带宽、分配、传输、同步还是存储。
• 记录硬件、包版本、dtype、形状、批大小和预热策略，并附上结果。

2. 移植前检查适用性

当热路径暴露出大量独立工作、运行频率足以摊销初始化和传输开销，并且工作集能够放入可用设备 内存且有余量容纳临时数据时，GPU 执行是有前景的。当工作负载较小、大部分是顺序执行、 主要由不受支持的操作主导，或需要频繁的主机-设备往返时，保留 CPU 路径。

不要将固定的行数阈值作为依据。对用户的实际形状和硬件进行基准测试。 对于超出内存的数据，在分配之前估算峰值工作内存，并选择分块、Dask 或流式设计。

3. 尝试影响最小的实现方式

1. 如果代码已经使用 GPU 原生框架，就在该框架内优化。
2. 尝试加速器或后端模式（cudf.pandas、cuml.accel、nx-cugraph）。
3. 仅在加速器覆盖范围或性能不足时，才迁移到原生 GPU API。
4. 仅当性能分析显示某项操作没有合适的库实现时，才编写自定义内核。

编写代码前，请先阅读相关库的参考文档；名称相同也可能在默认值、数据类型、输出类型和支持的参数上有所不同。

4. 保持连贯的 GPU 数据路径

• 仅传输一次输入，并让中间结果保留在设备上。
• 重用内存分配；语义允许时，优先使用 out= 或原地形式。
• 将小型操作批处理；如果融合逐元素操作可以避免生成中间数组，就进行融合。
• 只有在性能分析表明传输重叠很重要时，才使用固定页主机内存和非默认流。
• 仅当契约允许时，才选择 float32、混合精度或低精度存储。

5. 先验证语义，再追求速度

• 使用小型确定性测试样例和有代表性的数据，对比 CPU 和 GPU 输出。
• 为浮点结果设置明确的容差，并测试边界情况、NaN、排序和数据类型。
• 对于近似最近邻索引，对照精确搜索报告 recall@k；不要将精确 CPU 算法与近似 GPU 算法视为等价并进行比较。
• 检查加速器警告和日志，确认是否回退到了 CPU。

6. 正确地对 GPU 代码进行基准测试

GPU 工作是异步的，因此在未同步的调用外使用 CPU 计时器，测量到的只是入队时间。请先预热上下文创建和 JIT 编译，然后使用 CUDA 事件或支持相应库的计时器：

from cupyx.profiler import benchmark

print(benchmark(gpu_function, (arg1, arg2), n_warmup=10, n_repeat=100))

在笔记本中使用 %gpu_timeit，使用 Nsight Systems（nsys）查看端到端时间线，使用 Nsight Compute（ncu）分析内核。报告经过同步的内核/区域耗时和实际端到端延迟；如果生产环境会产生传输和转换开销，也要将其计入。

7. 保留、修改或放弃移植方案

仅当 GPU 路径通过正确性检查，且能在有代表性的数据上改善用户关心的指标时，才予以保留。如果没有改善，请说明限制因素是问题规模、传输、未支持导致的回退、内存压力、启动粒度，还是算法本身。

重要说明

• 如果应用需要可移植性，请提供 CPU 回退方案；否则应尽早失败，并明确报告硬件和依赖项错误。
• 对照 CPU 结果测试数值正确性（由于操作顺序不同，GPU 浮点结果可能略有差异）
• GPU 内存有限——如果数据集大于 GPU 内存，请考虑分块，或使用 RAPIDS Dask 进行多 GPU 处理
• 对支持的零拷贝数据交换，优先使用 CUDA Array Interface 或 DLPack；但请核实设备、数据类型、连续性、所有权和流语义，不要假设每次转换都没有开销。

参考文件

编写任何 GPU 优化代码之前，请先阅读相关参考文件：

文件	何时阅读
references/cupy.md	用户有 NumPy/SciPy 代码，或需要在 GPU 上执行数组操作
references/numba.md	用户有现成的 Numba-CUDA 代码，或需要显式 SIMT 内核；注意迁移到 Numba-CUDA-MLIR 的路径
references/cudf.md	用户有 pandas 代码，或需要在 GPU 上执行 DataFrame 操作
references/cuml.md	用户有 scikit-learn 代码，或需要在 GPU 上进行机器学习训练、推理或预处理
references/cugraph.md	用户有 NetworkX 代码，或需要在 GPU 上进行图分析
references/warp.md	用户需要用于仿真、空间计算、网格/体积查询、可微编程或机器人技术的 GPU 内核；如需高级物理引擎，请使用 Newton
references/kvikio.md	用户需要高速文件与 GPU 之间的 IO、GPUDirect Storage、将 S3/HTTP 数据读入 GPU，或在 GPU 上使用 Zarr
references/cuxfilter.md	用户正在维护或明确要求使用 cuxfilter（即将停止维护——26.06 是最终版本）
references/cucim.md	用户有 scikit-image 代码，或需要在 GPU 上进行图像处理、数字病理学或 WSI 读取
references/cuvs.md	用户需要在 GPU 上进行向量搜索、最近邻搜索、相似性搜索或 RAG 检索
references/cuspatial.md	用户正在维护或明确要求使用 cuSpatial（已归档——冻结在 25.04 版本，且已与当前 RAPIDS 隔离）
references/raft.md	用户需要稀疏特征值求解器、设备内存管理或多 GPU 基元

编写代码之前，请先阅读相应的参考文件——其中包含每个库特有的详细 API 模式、优化技巧和注意事项。

相关技能

• dask：用于超出内存容量的 pandas/NumPy 工作流的分布式计算。