NVIDIA nvmath-python是一个旨在连接Python科学计算社区与NVIDIA CUDA-X数学库的开发工具。它让Python用户无需改变现有工作流程，即可享受CUDA-X在常见数学运算中带来的性能优势。根据所调用的API类型，运算既可在CPU上执行，也可在支持CUDA的GPU，乃至分布式多GPU、多节点系统上运行。
随着nvmath-python v1.0正式发布，本文将深入介绍该库的设计理念与独特功能——涵盖从单CPU或单GPU到多GPU、多节点规模的数学运算加速能力。nvmath-python是对CUDA及NVPL数学库的Python抽象封装，底层支持cuFFT、cuBLASLt、cuDSS、cuSPARSE、cuTENSOR、cuBLASMp等多个库。此外，该库还引入了一种全新的稀疏性处理方案——通用稀疏张量（UST），用户可通过领域特定语言自定义适合特定应用场景的稀疏格式，无需手动编写底层实现代码。
安装说明
安装含有复杂原生依赖的Python包往往耗时且容易出错。nvmath-python的安装过程简洁高效，并支持针对不同环境进行定制：
支持pip、conda、uv、pixi等多种包管理器；可选择通过包管理器的依赖解析系统安装全部依赖，也可仅安装最小依赖集，适用于CI/CD流水线或纯CPU环境；可按需选择CPU后端、设备API支持或分布式API；可自由选择NumPy、CuPy或PyTorch等配套数组库（也可同时安装全部）。
设计定位
与NumPy等通用数学库类似，nvmath-python实现了大量适用于工程和科学计算的核心数值运算。但它并非旨在替代通用数组库，也不提供索引、切片或归约等传统功能。
nvmath-python的核心目标是将CUDA-X数学库的完整功能和强大性能引入Python生态，帮助现有数组库和框架更便捷地调用高度优化的GPU加速运算，从而无需依赖底层C/C++接口。
以下示例展示了nvmath-python如何接受NumPy数组作为输入，并返回NumPy数组作为结果：
import numpy as np
import nvmath
m, n, k = 10, 40, 100
a = np.random.randn(m, k)
b = np.random.randn(k, n)
c = nvmath.linalg.advanced.matmul(a, b)
多执行空间支持
nvmath-python支持灵活选择数组库，无论是GPU端的CuPy还是CPU端的NumPy均可兼容。这得益于以下底层支持：
GPU库方面，包含cuBLAS和cuFFT；CPU库方面，支持适用于NVIDIA Grace或ARM v8处理器的NVPL，以及适用于x86主机的Intel MKL；分布式库方面，支持cuBLASMp、cuSOLVERMp和cuFFTMp。
这一特性极大简化了CPU与GPU之间的代码迁移，并支持结合CPU与GPU执行的混合与分布式工作流。以下代码说明了nvmath-python如何同时支持多种内存和执行空间：
import cupy as cp
import numpy as np
import nvmath
N = 2048
a_gpu = cp.random.randn(N) + 1j * cp.random.randn(N)
a_cpu = np.random.randn(N) + 1j * np.random.randn(N)
c_gpu = nvmath.fft.fft(a_gpu)
c_cpu = nvmath.fft.fft(a_cpu)
每次调用的FFT执行空间会根据输入张量（a_gpu或a_cpu）自动推断，当然用户也可以手动指定执行空间。
通用API与专用API
nvmath-python的API大致分为两类：通用API和专用API。
通用API提供统一的用户体验，兼容多种执行空间、内存空间和操作数类型，类似于功能广泛的多用途工具，但可配置项仅限于各类场景共同支持的基础功能。专用API则针对特定操作范围提供全面的配置选项，能够充分发挥硬件性能，可能受限于特定硬件平台。
以矩阵乘法为例：advanced模块中的高级矩阵乘法专门针对GPU上的稠密操作数实现了D=f(AB+C)这一复合运算，提供充分挖掘硬件效率所需的全部配置选项；而通用矩阵乘法API则兼容CPU和GPU上的稠密与结构化操作数，但仅提供适用于更广泛场景的通用配置子集。
最优选择完全取决于具体使用场景：当某项运算成为计算瓶颈、需要针对特定硬件进行优化时，专用API是理想之选；而对于非性能关键型任务，或无需深度定制的场景，通用API更为适合。所有专用API均位于advanced子模块中，以便与通用API明确区分。
日志记录与数据流追踪
该库与Python标准库中的logging模块完整集成，支持在调试、信息、警告和错误等多个层级记录计算细节。
通过日志输出，用户可以直观了解操作数在内存与执行空间之间的流转情况。例如，当操作数来自CPU内存但在GPU上执行时，日志会明确提示可能存在的数据传输开销。nvmath-python会在可能的情况下自动选择最优执行空间，以最大限度降低数据传输代价，用户也可通过execution关键字参数手动指定所需的执行空间。
内核融合与性能提升
以链式调用方式实现D=f(αA·B+βC)这类运算时，若底层原语的算术强度较低，效率会明显不足。一个典型案例是对高瘦矩阵执行GEMM运算：
D=αA·B+βC
以下代码对比了CuPy与nvmath-python在高瘦矩阵上执行GEMM的方式：
import cupy as cp
import nvmath
m, n, k = 10_000_000, 40, 10
a = cp.random.randn(m, k, dtype=cp.float32)
b = cp.random.randn(k, n, dtype=cp.float32)
c = cp.random.randn(m, n, dtype=cp.float32)
alpha, beta = 1.5, 0.5
d1 = alpha * cp.matmul(a, b) + beta * c # 多个内核
d2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta) # 单个内核
融合的复合运算相比NumPy风格的API在性能上具有显著优势。nvmath-python底层依托cuBLASLt库，支持即时内核融合，是提升算术强度的有效手段之一。
有状态API与自动调优
上述示例均采用了nvmath-python的函数式（无状态）API，这是一种便捷的单次调用接口，但其内部包含耗时的准备逻辑，称为规划阶段。此外，准备阶段还可能包含自动调优的开销，与实际执行数学运算的执行阶段相互独立。
性能说明：NVIDIA CUDA-X数学库通过启发式算法来选择最优实现方案。由于针对特定问题规模、数据布局或数据类型优化的专用内核存在多个候选项，自动调优会迭代遍历各内核选项、测量其性能并选出最佳方案，因此该阶段可能较为耗时。
在深度学习等需要对不同输入重复执行相同操作的工作负载中，通过创建并复用计划对象，可以将规划开销分摊到多次执行中。nvmath-python的基于类（有状态）API正是为此场景而设计。
以下示例展示了如何使用有状态API进行矩阵乘法（含RELU_BIAS epilog），在一批矩阵上分阶段执行规划、自动调优和计算：
import nvmath
from nvmath.linalg.advanced import MatmulEpilog
import cupy as cp
feed_count = 10
batch_size = 1024
m, n, k = 1024, 1024, 1024
a = cp.random.rand(batch_size, m, k, dtype=cp.float32)
b = cp.random.rand(batch_size, k, n, dtype=cp.float32)
bias = cp.random.rand(batch_size, m, 1, dtype=cp.float32)
with nvmath.linalg.advanced.Matmul(a, b) as mm:
mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),
epilog_inputs={"bias": bias})
mm.autotune(iterations=5)
for i in range(feed_count):
d = mm.execute()
mm.reset_operands_unchecked(a=d)
与无状态API相比，有状态API通过分摊规划和准备开销显著降低了单次执行的总成本，而自动调优则能在此基础上进一步提升性能。自动调优结果可序列化保存到磁盘，并在新会话中加载复用，避免重复调优。测试表明，内置启发式算法在多数情况下已能选出高性能内核；但对于某些问题规模、数据类型、操作数布局与硬件的特定组合，自动调优仍可带来额外增益。
与Python编译器集成
nvmath-python与numba-cuda等Python编译器深度集成，支持将高性能自定义Python代码即时编译（JIT），并与nvmath-python的运算协同使用。
FFT回调示例
FFT的回调函数以预定义签名的Python函数形式编写，经JIT编译为中间表示后，作为nvmath-python正向或反向FFT的自定义前处理或后处理逻辑。
以高斯滤波器为例，该滤波器对图像进行模糊处理。实现流程为：原始图像 → R2C FFT → 高斯滤波 → C2R iFFT → 滤波结果图像。高斯滤波器的频域形式为H(fx,fy)=exp(-2π?σ?(fx?+fy?))，以下代码展示了具体实现：
from PIL import Image
import nvmath
import cupy as cp
img = cp.asarray(Image.open("your_lovely_dog.jpg").convert("L")) / 255.0
wh = img.shape[0] * img.shape[1]
sigma_value = 20.0
def gaussian_filter(shape, sigma):
fy = cp.fft.fftfreq(shape[0])[:,None]
fx = cp.fft.rfftfreq(shape[1])[None,:]
return cp.exp(-2.0 * cp.pi * cp.pi * sigma * sigma * (fx * fx + fy * fy))
def epilog_impl(data_out, offset, data, filter_data, unused):
data_out[offset] = data * filter_data[offset] / wh
epilog = nvmath.fft.compile_epilog(epilog_impl, "complex64", "complex64")
h_filter = gaussian_filter(img.shape, sigma_value)
img_fft = nvmath.fft.rfft(img, epilog={"ltoir": epilog, "data": h_filter.data.ptr})
filtered_img = nvmath.fft.irfft(img_fft)
设备API与Monte Carlo仿真
nvmath-python还支持在numba-cuda编写的GPU内核中调用设备API，涵盖FFT、GEMM、稠密直接求解器（LU、Cholesky、QR）和随机数生成器（RNG）。以下示例实现了用于Monte Carlo股价仿真的几何布朗运动（GBM），采用nvmath-python的高斯分布随机数生成器，并结合自定义numba-cuda代码将正态分布转换为GBM路径：
from numba import cuda
from nvmath.device import random
import cupy as cp
import math
compiled_rng = random.Compile(cc=None)
rng_seed = 7777
n_time_steps, n_paths = 252, 8192
mu, sigma, s0 = 0.003, 0.027, 100.0
threads_per_block = 32
blocks = n_paths // threads_per_block + bool(n_paths % threads_per_block)
nthreads = threads_per_block * blocks
@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)
def init_rng(states, seed):
idx = cuda.grid(1)
random.init(seed, idx, 0, states[idx])
@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)
def generate_gbm_paths(states, paths, nsteps, mu, sigma, s0):
idx = cuda.grid(1)
if idx >= paths.shape[0]:
return
paths[idx, 0] = s0
for i in range(1, nsteps, 4):
v = random.normal4(states[idx])
vals = v.x, v.y, v.z, v.w
for j in range(i, min(i + 4, nsteps)):
paths[idx, j] = paths[idx, j-1] * math.exp(mu + sigma * vals[j-i])
states = random.StatesPhilox4_32_10(nthreads)
init_rng[blocks, threads_per_block](states, rng_seed)
paths = cp.empty((n_paths, n_time_steps), dtype=cp.float32, order='F')
generate_gbm_paths[blocks, threads_per_block](states, paths, n_time_steps, mu, sigma, s0)
快速上手
nvmath-python在兼顾开发效率的同时，不以牺牲性能为代价，代表了现代数学库设计理念的全新探索。只需一条命令即可开始使用：
pip install nvmath-python[cu13]
更多学习资源包括：详细安装文档、包含代码示例和深度教程的nvmath-python GitHub仓库、NVIDIA加速计算中心的Python进阶训练材料，以及关于通用稀疏张量（UST）的专题文章。
该库凝聚了NVIDIA众多团队成员的心血，核心贡献者包括：Harun Bayraktar、Becca Zandstein、Lukasz Ligowski等数十位工程师。
Q&A
Q1：nvmath-python是什么？和NumPy有什么区别？
A：nvmath-python是NVIDIA推出的Python数学运算库，用于将CUDA-X高性能数学库（如cuFFT、cuBLASLt等）暴露给Python用户。与NumPy不同，nvmath-python不提供索引、切片或归约等通用数组操作，而是专注于在CPU、单GPU乃至多GPU分布式环境下提供高度优化的数学运算加速，并且支持NumPy、CuPy、PyTorch等多种数组库作为输入和输出。
Q2：nvmath-python的有状态API和无状态API有什么区别？什么时候该用哪种？
A：无状态API是单次调用接口，使用方便但每次调用都会执行规划和准备阶段，开销较大。有状态API支持将规划、自动调优和执行拆分为独立阶段，适合需要对不同输入重复执行相同操作的场景（如深度学习训练），可将规划开销分摊到多次执行中。自动调优结果还可序列化保存和跨会话复用。如果某项运算不是性能瓶颈，用无状态API更简洁；如果是高频重复运算，推荐使用有状态API。
Q3：nvmath-python如何安装？支持哪些环境？
A：只需一条命令即可安装：pip install nvmath-python[cu13]。该库支持pip、conda、uv、pixi等多种包管理器，可选择完整依赖安装或最小化安装（适合CI/CD或纯CPU环境），支持选择CPU后端、GPU设备API或分布式API，并兼容NumPy、CuPy、PyTorch等主流数组库。详细安装选项可参考官方安装文档。
