Deep Dive into Math Libraries

Arthy Sundaram, Product Manager
Harun Bayraktar, Director of Engineering
GTC Spring 2024

引言

数学库是实现满足您需求的加速的基础。核心优势包括：

无缝加速 (Seamless Acceleration): 利用张量核心 (Tensor Cores) 等硬件特性进行加速。
可组合性 (Composability): 提供设备扩展库 (Device Extension Libraries)，增强灵活性。
NVPL on Arm: 为 Arm 架构提供高性能 CPU 库。
多GPU与多节点 (Multi-GPU & Multi-Node): 能够扩展至超级计算机规模。

议程 (Agenda)

本次演讲将涵盖以下主题：

快速傅里叶变换库 (Fast Fourier Transform Libraries)
稠密和稀疏线性代数库 (Dense and Sparse Linear Algebra Libraries)
高级求解器 (Advanced Solvers)
图像和数据压缩库 (Image and Data Compression Libraries)
CPU 库 (CPU Libraries)
Python API 介绍 (Introducing Python APIs)
结束语 (Closing Remarks)

快速傅里叶变换库 (Fast Fourier Transform Libraries)

cuFFT API 扩展

cuFFT 提供了多种 API 扩展，以满足不同场景的需求，这些 API 均可从 CPU 代码中调用。

API 扩展	描述	可用性
cuFFT Host APIs	用于单 GPU 的 GPU 加速 FFT	CUDA Toolkit/HPC SDK
cuFFTxt Host APIs	多 GPU 单节点 FFT	CUDA Toolkit/HPC SDK
cuFFTMp Host APIs	多 GPU 多节点 FFT	HPC SDK
cuFFTDx Device APIs	用于在 CUDA 核函数内部执行的 FFT (In-kernel FFTs)	独立下载 (Standalone download)

cuFFT: 即时编译与链接时优化 (Just-In-Time, Link-Time Optimization)

传统方法为每种可能的配置和 GPU 架构提供预编译的专用核函数，但这会导致“组合爆炸”问题，扩展性差。

JIT LTO (即时编译与链接时优化) 通过将核函数作为预编译的“片段”提供，解决了这一问题。这些片段在运行时（Runtime）与用户代码（例如回调、卷积等）结合，动态生成针对目标方案的专用核函数。这为用户代码集成提供了更多可能性，并使库功能更加丰富。

cuFFT - JIT LTO 性能与路线图

JIT LTO 实现了核函数的运行时最终确定和内联回调，并增加了对 Windows 的支持。

CUDA 12.4: 提供了性能更强的 64 位核函数。在 H100 上的测试显示，与 CUDA 12.3 相比，对于一系列 FFT 尺寸，性能有显著提升，最高可达 5 倍以上。
公开早期访问 (In Public Early Access): 提供了性能更强的 LTO 回调函数并支持 Windows。在 H100 上的 FP32 C2C 测试显示，LTO 回调函数相比 cuFFT 11.7 有明显性能提升。
未来规划 (Coming to cuFFT): 将支持 R2C、C2R 及其他类型的核函数，并提供归一化、零填充和混合精度支持。

cuFFTDx

cuFFTDx 是一个 C++ 头文件库，提供设备端 API，用于将 FFT 融合（fusing）到用户核函数中。

核心特性:
- 这是一个 C++ 头文件库，用于在用户核函数中内联 FFT。
- 当数据能够放入共享内存或寄存器时，性能表现最佳。
- 可以在单个核函数中将 FFT 与自定义用户操作或其他 Dx API 融合。
性能优势:
- 通过将多个操作（如加载、FFT、用户代码、存储）融合到单个 CUDA Kernel 中，cuFFTDx 减少了对全局内存的读写，显著提升了性能。
- 如下图所示，在 H100 上进行一维卷积（FP32）测试，使用 cuFFTDx 的融合核函数（1 fused kernel）相比使用 cuFFT 的三个独立核函数（3 Kernels）性能有数倍提升。

稠密和稀疏线性代数库 (Dense and Sparse Linear Algebra Libraries)

cuBLAS

cuBLAS 是用于执行基本线性代数运算（BLAS）的 GPU 加速库。

cuBLAS API 扩展

cuBLAS 提供了多种可从 CPU 代码调用的 API 扩展。

API 扩展	描述	可用性
cuBLAS Host APIs	标准的 GPU BLAS API，带有额外的 GEMM 扩展。	CUDA Toolkit/HPC SDK
cuBLASLt Host APIs	用于 GEMM 的高级 API，具备可编程性、AI 训练的启发式核函数选择以及多种混合精度和尾声（epilogues）支持。	CUDA Toolkit/HPC SDK
cuBLASXt Host APIs	用于单节点多 GPU 的 cuBLAS 调用。	CUDA Toolkit/HPC SDK

新的 API 扩展: cuBLASDx 和 cuBLASMp

这些新的 API 扩展现已开放下载。

API 扩展	描述	可用性
cuBLASMp Host APIs (Preview)	用于多 GPU 多节点（MGMN）的标准 BLAS API，可从 CPU 代码调用。	独立下载和 HPC SDK
cuBLASDx Device side APIs (Preview)	简洁易用的在核函数内部使用的 GEMM（in-kernel GEMMs）。	独立下载

cuBLASMp: 一个高性能、多进程的 GPU 加速库，用于分布式稠密线性代数。
cuBLASDx: cuBLAS 的设备端 API 扩展，用于在 CUDA 核函数内部执行 BLAS 计算。融合数值操作可以减少延迟并提高应用程序性能。

cuBLASDx: 用于核函数内 GEMM 的设备端 API

cuBLASDx 旨在简化在 CUDA 核函数内部执行通用矩阵乘法（GEMM）的操作。

核心特性:
- 构建于 CuTE 之上，为 CUDA 核函数内的 GEMM 提供了更简洁的 API。
- 支持内联用户代码和 cuFFTDx 调用。
- 支持所有能放入共享内存的矩阵尺寸，包括实数、复数、FP16、FP32 和 FP64。
未来规划: 将支持混合精度和窄精度格式。
性能: 通过将多个 GEMM 和自定义操作融合，cuBLASDx 能实现更高的 GPU 利用率。如下图所示，在 H100 上的测试表明，融合多个操作后，cuBLASDx 的性能相比 cuBLAS+Thrust 有显著提升（最高可达 3 倍）。

cuBLASMp: 分布式稠密线性代数

cuBLASMp 提供多节点主机 API 以实现规模化计算。

特点:
- 支持 GEMM, SYRK, TRSM 等操作。
- 每个 GPU 一个进程。
- 支持 FP16, FP32, FP64。
未来规划: 将支持混合精度和窄精度格式。
性能: 在 DGX-H100 集群上进行的 cuBLASMp PDGEMM 弱扩展性测试显示，随着 GPU 数量增加到 1024 个，其性能可以达到理论峰值的 66%。测试中的问题规模为每个 GPU M=N=K=55k。

cuBLAS API 扩展家族概览

下图总结了 cuBLAS API 扩展的完整家族，涵盖了从单 GPU 到多节点集群，以及从主机端到设备端的各种应用场景。

cuBLAS - Grouped GEMM

从 CUDA 12.4 开始引入 Grouped GEMM (实验性功能)。这是一个单一、简化、高性能的 API，用于对不同大小和 leading dimensions 的矩阵进行批处理。

应用场景包括：
- 专家混合（Mixture of Experts）: 激活较小模型的子集。
- 异构图神经网络（Heterogeneous GNN）: 学习层次关系。

这些应用场景会产生大量来自不同专家或不同子图的相异 GEMM（通用矩阵乘法）操作，Grouped GEMM 可以将它们分组处理以提升效率。

Image illustrating the concept of Grouped GEMM for Mixture of Experts and Heterogeneous GNN models.

cuBLAS - Grouped GEMM vs 迭代式批处理 GEMM

Grouped GEMM 提供单一、简化、高性能的 API，用于对不同大小和 leading dimensions 的矩阵进行批处理。

优势:
- 单次内核提交: 无启动延迟，提升 GPU 利用率。
- 支持的数据类型: 支持 FP32、TF32 和 FP64。
- 即将支持: FP16 和 BF16。

性能对比:
下图展示了在 H100 SXM 上，专家混合（MoE）生成阶段 GEMM 的性能。结果显示，与迭代式批处理 GEMM 相比，cuBLAS Grouped GEMM 实现了显著的几何平均加速。
- 批大小为 64 时，加速比为 1.2 倍。
- 批大小为 8 时，加速比为 1.18 倍。

Comparison of Grouped Batched GEMMs (concurrent kernels) vs Iterative calls to Batched GEMMs (sequential kernels) and performance benchmark graph.

cuBLASLt

功能覆盖与性能提升

cuBLASLt 提供了用于 GEMM 的高级 API，支持多种混合精度、epilogues、内核选择和性能优化，是 CUDA Toolkit/HPC SDK 的一部分。

近期更新:
- 增加了对更大矩阵维度（m, n）和批大小的支持。
- 改进了针对最新大语言模型（LLM）的 matmul 性能。

性能基准测试:
- MLPerf v3.0 GPT3 LLM 训练 (仅 Matmul 加速比):
- 在 bf16 精度下，H100-SXM 和 H200-SXM 相较于 A100 实现了 3.1 倍的加速。
- 在 bf16+fp8 混合精度下，H100-SXM 和 H200-SXM 分别实现了 4.6 倍和 4.8 倍的加速。

Llama2 70B 训练 (仅 Matmul 加速比):
- 在 bf16 精度下，H100-SXM 和 H200-SXM 相较于 A100 分别实现了 2.7 倍和 2.8 倍的加速。

Performance charts showing matmul speedups for MLPerf v3.0 GPT3 and Llama2 70B training on different GPUs.

cuTENSOR: GPU 加速的张量线性代数库

2.0 版本的新 API、功能和性能

API 一致性: 基于 Plan 的多阶段 API 扩展至所有操作。
动态分配描述符: 支持任意维度的张量描述符。
“可选加入”的即时编译（Just-in-Time）: 为张量缩并（tensor contractions）生成优化内核。
获取方式: 可通过 https://developer.nvidia.com/hpc-sdk 或 https://developer.nvidia.com/cutensor/downloads 获取。
生态集成: 已被 cuQuantum、CuPy、CUDA Julia 等库采用。

性能提升:
- cuTENSOR 2.0.0 (无 JIT) vs 1.7.0 加速比: 在 A100-80GB-PCIe 和 H100-80GB-PCIe 上，新版本在 bf16、fp16、fp32、fp64 和 complex(fp32) 数据类型上均有性能提升，最高加速比接近 3 倍。
- cuTENSOR 2.0 - JIT 带来的增量加速 (H100): 在 H100 上，启用 JIT 相比未启用 JIT，可带来额外的性能增益，最高可达 1.4 倍以上。

Key features of cuTENSOR 2.0 and performance comparison charts.

cuSPARSE

稀疏性实践：结构化与非结构化

cuSPARSE 库提供了处理不同类型稀疏性的 API：

cuSPARSE Host APIs:
- 非结构化稀疏: 通过单一 API 调用实现标准的稀疏 BLAS 计算。适用于元素随机分布的稀疏矩阵。
cuSPARSELt Host APIs:
- 细粒度 2:4 结构化稀疏: 通过无状态、多阶段的 API 访问稀疏张量核心（Sparse Tensor Cores）。适用于矩阵中每4个元素有2个非零元素的结构化稀疏模式，可实现剪枝和压缩。

Diagram illustrating unstructured sparsity and the process of fine-grained 2:4 structured sparsity.

cuSPARSE: 新的稀疏格式 - 分块压缩稀疏行格式 (BCSR)

BSR/BCSR 支持 SDDMM (采样密集-密集矩阵乘法): 适用于具有密集块子结构的矩阵。
业界评价: "BCSR 格式可以增强 PaddlePaddle 中稀疏训练和推理的性能" - 百度。

性能对比:
下图展示了 BSR SDDMM 与 CSR SDDMM 的性能对比 (单位: GFLOPS)。结果表明，随着块大小 (blocksize) 的增加，BSR SDDMM 的性能显著优于 CSR SDDMM。例如，当块大小为 32 时，BSR SDDMM 的性能达到 1,354.88 GFLOPS，而 CSR SDDMM 仅为 501.2 GFLOPS。

Performance graph comparing BSR SDDMM versus CSR SDDMM across different block sizes.

cuSPARSELt 0.6.0

该版本增加了对 Hopper Tensor Cores 的支持，用于矩阵的 2:4 结构化稀疏计算。

Hopper 支持: 需要 R535 TRD7、R550 TRD1 或更新版本的驱动程序。
混合精度支持: 支持 FP8 (e4m3 和 e5m2) 及其他窄数据类型。
其他支持的数据类型: INT8、FP16、BF16、TF32。
未来规划: 将支持更多的矩阵尺寸和精度类型。
生态集成: cuSPARSELt + Ampere (0.5.x 版本) 现已通过 PyTorch 2.1+ 提供。

性能:
下图展示了在 H100 上，使用 cuSPARSELt 的 2:4 稀疏 GEMM 相对于使用 cuBLASLt 的密集 GEMM 的加速比。在不同矩阵尺寸和数据类型（bf16, int8, fp8）下，均可获得约 1.5 倍的加速。

Performance chart showing speedup of 2:4 sparsity using cusparselt vs dense GEMM on H100.

高级求解器 (Advanced Solvers)

cuDSS: 适用于 NVIDIA GPU 的直接稀疏求解器

cuDSS 0.2.0 版本现已在 https://developer.nvidia.com/cudss 上提供。

求解流程:
1. 分析/重排序 (ANALYSIS/REORDERING):
- 矩阵类型: 实数/复数、对称/非对称、厄米矩阵。
- 重排序算法: AMD, COLAMD, RTF, Metis ND。

因子分解 (FACTORIZATION):
- 分解方法: LU, Cholesky, Bunch-Kaufman。
- 主元选择 (Pivoting): 无、全局、局部、列、行、阈值。
求解 (SOLVING):
- 支持多个右端项 (RHS)、转置求解。
- 可选的迭代求精以提高精度。

平台支持:

支持 Windows x86、Linux x86 和 Arm (Grace CPU)。
多节点多 GPU 支持正在规划中。
应用案例: Honeywell 使用 cuDSS 加速 UniSim 设计过程仿真。

Flowchart of the cuDSS direct sparse solver process.

图像和数据压缩库 (Image and Data Compression Libraries)

nvImageCodec: GPU 加速的图像处理

统一 API: 提供用于编码/解码图像的统一 C 和 Python API。
开源: 采用 Apache 许可证。
异构批处理: 支持对不同类型的图像进行批处理。
第三方编解码器支持: 支持第三方编解码器，并可设置优先级和回退机制。
格式支持: 支持最流行的图像格式，如 jpeg, jpeg2000, htj2k, tiff, bmp, png, pnm, webp。
平台支持: x86_64 (Linux, WSL), arm64_sbsa (Linux)。
版本: V0.2.0 已可下载 (https://developer.nvidia.com/nvimgcodec, https://pypi.org/project/nvimgcodec/, https://github.com/NVIDIA/nvImageCodec)。

nvImageCodec 支持矩阵:
下表详细列出了不同格式在解码和编码时所支持的后端库。

nvImageCodec Support Matrix table showing supported formats and backends for decoding and encoding.

nvCOMP: GPU 加速的数据压缩

压缩算法: 支持 zSTD、deflate、gDeflate、LZ4、snappy、bitcomp、cascaded、gANS 等。
集成:
- Spark RAPIDS (cuDF): [端到端查询加速高达 10 倍]
- Parabricks (医疗保健): [fq2bam 加速 30%]
- nvTIFF: [支持 deflate]
获取方式: 可从 https://developer.nvidia.com/nvcomp 下载。

H100 性能:
下图展示了不同算法在 H100 上的压缩率、压缩吞吐量和解压吞吐量。

Performance charts for nvCOMP on H100, showing Compression Ratios, Compression Throughputs, and Decompression Throughputs for various algorithms.

CPU 库 (CPU Libraries)

NVPL-24.03 Beta：支持 NVIDIA Grace 上 HPC 应用的 CPU 数学库

NVPL (NVIDIA Performance Libraries) 是一套用于 NVIDIA Grace CPU 架构的高性能计算（HPC）数学库。

NVPL-24.03 现已发布
- 下载地址：http://developer.nvidia.com/nvpl-downloads
- 文档：http://docs.nvidia.com/nvpl
- 示例：http://github.com/NVIDIA/NVPLSamples
引入 NVPL Tensor
- API 基于 cuTENSOR 2.0
- 支持高达 64 维的张量
自 23.11 版本以来的其他改进
- NVPL BLAS：小型 DGEMM 优化；针对特定 HPC 应用的端到端（E2E）性能提升 5%。
- NVPL FFT：复数到复数（C2C）性能优化；针对大数据集的多线程支持。
- NVPL Sparse：稀疏向量求解（SpSV）速度提升 1.8 倍，内存占用减少 1.8 倍。

Page 31: NVPL-24.03 Beta 版本发布详情、性能图表和 Grace Superchip 架构。

上图右侧的性能图表显示：

端到端 Grace 上的加速（BLAS 部分）：与 NVPL 基础版和 OpenBLAS 相比，NVPL 24.03 在 NWChem、CP2K 和 Arpege 等应用中均实现了显著的性能提升。
1D C2C FP64 FFT 性能：在 Grace Superchip 上，使用 144 线程、批处理大小为 1 的情况下，NVPL FFT 相较于 FFTW 实现了高达 200 倍以上的加速。

下图展示了 NVIDIA Grace Superchip 的架构，其包含 144 核 Arm Neoverse V2 核心，通过 NVLink-C2C 连接，提供 900 GB/s 的双向带宽和 1TB/s 的 LPDDR5X ECC 内存带宽。

Python API 介绍 (Introducing Python APIs)

本节内容由 Harun Bayraktar, Leo Fang, Satya Varadhan, Leopold Cambier 呈现。

Python 数学库功能的生态系统

在过去十年中，Python 已成为最流行的编程语言，主要由数据科学和人工智能驱动，同时在科学计算领域也取得了长足的进步。
GPU 加速的数学库功能可通过流行的框架和库（如 CuPy, JAX, PyTorch, Numba）获得，但功能覆盖范围是零散的，且性能并非总是最优。这些高级框架构建在 CUDA 数学 API（如 CUBLAS, CUSPARSE, CUTENSOR 等）之上。

探路者：cuQuantum Python

cuQuantum Python 是为张量网络收缩提供 Pythonic 用户体验（UX）的成功案例。它展示了高级 API 如何显著简化代码并提高开发效率。

C++ 示例: 279 行代码 (SLOC) + 61 行注释
低级 Python 绑定示例: 148 行代码 (SLOC) + 60 行注释
Pythonic API 示例: 仅需几行高级代码即可完成复杂的张量收缩操作。

Page 34: 从 C++ 到低级 Python 绑定再到 Pythonic API 的代码行数对比。

nvmath-python 架构

nvmath-python 旨在以便捷的 Pythonic 方式高性能地访问 CUDA 数学库功能。

核心特性：
- 最小化延迟：从 Python 快速访问新的 CUDA 库功能。
- 互操作性：与核心 Python 数值计算包（NumPy, SciPy, CuPy, PyTorch 等）兼容。
- Pythonic 库接口：为 CUDA 加速库提供 Pythonic API，这些 API 构建在低级 1:1 C API 的 Python 绑定之上。
- 平台无关性：可在 GPU、CPU (x86 & Arm) 和 Grace-Hopper 上运行。

设计优势：
- 为终端用户、包开发者和内核作者提供统一的高生产力与高性能设计。
- 能够访问所有扩展和峰值性能功能，如降低和混合精度、设备 API、以及功能定制（例如 cuFFT 内联回调）。

Page 35: nvmath-python 的架构图，展示了其如何连接 Python 应用与底层硬件加速库。

nvmath-python 演示

示例 1：专用矩阵乘法（matmul）API 的基本和高级用法
- 操作: d = ReLU(a x b + bias)
- 这是一个融合操作的例子，详情请参见 S62162_nvmath-python_fused_matmul_demo.pdf。

示例 2：三种方式实现卷积
- 操作: y = iFFT(FFT(xin) . f)
- 实现方式包括：主机 API、回调函数、融合内核。
- 详情请参见 S62162_nvmath-python_convolutions_demo.pdf。

nvmath-python 演示总结

nvmath-python 提供了一个高生产力的环境，同时赋予了优化性能的灵活性。

示例 1 性能（矩阵乘法）：使用 cuBLASLt 的专用 matmul API。与分成 3 个独立内核相比，自动调优和融合内核（Algo 0）能够达到更高的峰值 TFLOPS 百分比，性能分别达到 89.9% 和 84.2%。
示例 2 性能（卷积）：使用 cuFFT 和 cuFFTDx。与使用 3 个独立内核的主机 API 相比，使用主机 API 加回调函数（2个内核）可获得 1.28 倍的加速，而使用设备 API（融合内核）则可实现 2.25 倍的加速。

发布时间线

nvmath-python 的发布将通过 pip 和 conda 提供，源代码将在 GitHub 上开源。

2024 H1 (上半年)
- 支持回调的 FFT
- 密集线性代数
- 设备扩展 (Numba)
- 仅支持 GPU
2024 H2 (下半年)
- 为设备扩展支持更多编译器 (如 Warp)
- 稀疏线性代数
- 随机数生成器 (RNG)
- 支持 CPU (NVPL) 和 GPU
未来
- 多 GPU (MG) 和多 GPU 多节点 (MGMN) 的 FFT、线性代数等。

结束语 (Closing Remarks)

NVIDIA 数学库总结

NVIDIA 数学库在以下四个关键领域持续发展：

无缝加速 (Seamless Acceleration)：利用 Tensor Cores 等硬件特性。
- 在 cuBLASLt 和 cuSPARSELt 中支持 FP8
- 涵盖 cuBLAS, cuSPARSE, cuDSS, cuTENSOR, cuFFT, nvImageCodec, nvCOMP
- nvmath-python
可组合性 (Composability)：通过设备扩展库实现。
- cuFFTDx
- cuBLASDx
- Python 设备扩展
NVPL on Arm：高性能 CPU 库。
- NVPL 24.03 发布，新增 Tensor 库。
多 GPU 与多节点 (Multi-GPU & Multi-Node)：扩展至超级计算机。
- cuBLASMp 加入 cuSOLVERMp 和 cuFFTMp MGMN 库系列。

参考文献

S61198 CUTLASS: A Performant, Flexible, and Portable Way to Target Hopper Tensor Cores
S63345 Warp: Advancing Simulation AI with Differentiable GPU Computing in Python
S61598 Scientific Computing With NVIDIA Grace and the Arm Software Ecosystem
S61203 A Deep Dive into the Latest HPC Software
S62515 GPU-Accelerating Process Simulation Performance using NVIDIA's cuDSS Sparse Linear Systems Solver
S62275 Performance Optimization for Grace CPU Superchip
S62400 CUDA: New Features and Beyond
CWE62240 NVIDIA Math Libraries