optimize-for-gpu

用途

该技能面向在 NVIDIA 硬件上加速科学计算 Python 代码的需求，核心原则是把 GPU 加速当作以证据驱动的优化，而不是自动重写。它要求保留用户原有的数值与算法契约，用有代表性的数据测量，并且只有在同步的端到端基准测试显示有实际提升时，才保留 GPU 版本。

适用场景

技能适用于以下情况：用户希望加速数值/科学 Python 代码；处理大型数组、矩阵或数据帧；提到 CUDA、GPU、NVIDIA 或并行计算；使用 NumPy、pandas、SciPy、scikit-learn、NetworkX 或 scipy.sparse.linalg 处理大数据集；需要底层 GPU 原语（稀疏特征求解器、设备内存管理、多 GPU 通信）；进行机器学习（训练、推理、超参数调优、预处理）；进行图分析（中心性、社区发现、最短路径、PageRank 等）；进行向量搜索、最近邻搜索、相似度搜索或构建 RAG 流水线；使用 Faiss、Annoy、ScaNN 或 sklearn NearestNeighbors；需要 GPU 加速的交互式仪表板、交叉过滤或大数据集探索性分析；使用 GeoPandas 或 shapely 进行地理空间分析；使用 scikit-image 或 OpenCV 进行图像处理、计算机视觉或医学影像；处理全切片图像（WSI）、数字病理、显微镜或遥感影像；将大型二进制数据文件加载到 GPU 内存；从 S3、HTTP 或 WebHDFS 直接读取文件到 GPU 内存；提到 GPUDirect Storage（GDS）；进行物理仿真或可微分仿真；需要网格操作或几何处理；进行机器人学（运动学、动力学、控制）；有可 JIT 编译到 GPU 内核的 Python 仿真循环；提到 NVIDIA Warp；进行仿真、信号处理、金融建模、生物信息学、物理等计算密集型工作；希望优化现有代码且 GPU 加速是正确方向。

选择最小合适的加速层

优先使用维护良好的库实现，而非自定义内核。对应关系如下：NumPy/SciPy 对应 CuPy（数组、稀疏矩阵、线性代数、FFT、信号处理）；pandas 对应 cudf.pandas，其次 cuDF（先加速器模式，需要更多控制时用原生 API）；scikit-learn 对应 cuml.accel，其次 cuML；NetworkX 对应 nx-cugraph，其次 cuGraph；scikit-image 对应 cuCIM；Faiss/Annoy/k-NN 对应 cuVS；原始或远程文件 I/O 对应 KvikIO；自定义数组内核对应 Numba-CUDA-MLIR（新工作）或 Numba-CUDA（现有代码）；空间或可微分内核对应 Warp；高层物理仿真对应 Newton；底层 RAPIDS 原语对应 RAFT（pylibraft）。

不要仅仅为了使用这些库而把代码移出 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。

优化工作流

定义契约与基线：捕获有代表性的输入、期望输出和可接受的数值容差；测量当前端到端路径（输入、传输、计算、输出）；在改代码前做性能剖析，识别真正的瓶颈是计算、内存带宽、分配、传输、同步还是存储；记录硬件、包版本、dtype、形状、批大小和预热策略。
移植前检查适用性：当热路径有大量独立工作、运行足够频繁以摊销初始化和传输成本、工作集能放入设备内存并留出临时空间时，GPU 执行才有前景。工作负载小、大多串行、被不支持的操作主导或需要频繁主机-设备往返时，保留 CPU 路径。不要用固定行数阈值作为依据，应针对用户实际形状和硬件做基准测试。对于超出内存的数据，先估算峰值工作内存，再选择分块、Dask 或流式设计。
尝试侵入性最小的实现：如果代码已使用 GPU 原生框架，就在该框架内优化；尝试加速器或后端模式（cudf.pandas、cuml.accel、nx-cugraph）；仅在加速器覆盖或性能不足时转向原生 GPU API；仅当剖析显示某操作没有合适的库实现时才写自定义内核。写代码前阅读相关库参考，因为兼容的名称在默认值、dtype、输出类型和支持参数上仍可能不同。
保持一致的 GPU 数据路径：输入只传输一次，中间结果常驻设备；复用分配，在语义允许时优先使用 out= 或原地形式；批处理小操作，融合逐元素工作以消除中间数组；仅在剖析显示传输重叠重要时才使用固定主机内存和非默认流；仅在契约允许时选择 float32、混合精度或降低精度存储。
先验证语义再追求速度：在小规模确定性夹具和有代表性数据上比较 CPU 与 GPU 输出；对浮点结果使用显式容差，并测试边界情况、NaN、顺序和 dtype；对近似最近邻索引，报告相对精确搜索的 recall@k，不要把精确 CPU 算法与近似 GPU 算法当作等价比较；检查加速器警告和日志中的 CPU 回退。
正确基准测试 GPU 代码：GPU 工作是异步的，未同步调用外部的 CPU 计时器只测量入队时间。先预热上下文创建和 JIT 编译，然后使用 CUDA 事件或库感知计时器，例如：

from cupyx.profiler import benchmark

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

在 notebook 中使用 %gpu_timeit，用 Nsight Systems（nsys）看端到端时间线，用 Nsight Compute（ncu）做内核分析。同时报告同步后的内核/区域时间和真实的端到端延迟；当生产环境需要支付传输和转换成本时，应将其计入。

保留、修改或放弃移植：仅当 GPU 路径通过正确性检查，并在有代表性数据上改善了用户关心的指标时才保留。否则，解释限制因素是问题规模、传输、不支持的回退、内存压力、启动粒度还是算法本身。

重要说明

应用需要可移植性时提供 CPU 回退；否则以清晰的硬件和依赖错误尽早失败。针对 CPU 结果测试数值正确性（GPU 浮点可能因操作顺序略有不同）。GPU 内存有限，数据集大于 GPU 内存时考虑分块或使用 RAPIDS Dask 进行多 GPU。优先使用 CUDA Array Interface 或 DLPack 进行支持情况下的零拷贝交换，但要验证设备、dtype、连续性、所有权和流语义，不要假设每次转换都是免费的。

参考文件

在编写任何 GPU 优化代码前，阅读相关参考文件：references/cupy.md（NumPy/SciPy 代码或需要 GPU 数组操作）；references/numba.md（现有 Numba-CUDA 代码或需要显式 SIMT 内核，注意迁移到 Numba-CUDA-MLIR 的路径）；references/cudf.md（pandas 代码或需要 GPU 数据帧操作）；references/cuml.md（scikit-learn 代码或需要 GPU 上的 ML 训练/推理/预处理）；references/cugraph.md（NetworkX 代码或需要 GPU 图分析）；references/warp.md（需要用于仿真、空间计算、网格/体积查询、可微分编程或机器人学的 GPU 内核；高层物理引擎用 Newton）；references/kvikio.md（需要高性能文件 I/O 到/从 GPU、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 模式、优化技巧和陷阱。

技能介绍

用途