借助 NVIDIA nvmath-python 实现大规模高性能核心数学运算

NVIDIA nvmath-python 是一个旨在弥合 Python 科学计算社区与 NVIDIA CUDA-X 数学库之间差距的库。它为 Python 用户提供对 CUDA-X 性能的访问，以执行常见…

NVIDIA Technical Blog · Jul 30, 2026

Run High-Performance Core Math at Scale with NVIDIA nvmath-python

NVIDIA nvmath-python 是一个旨在弥合 Python 科学计算社区与 NVIDIA CUDA-X 数学库之间差距的库。它为 Python 用户提供对 CUDA-X 性能的访问，以执行常见数学运算，而无需中断现有工作流程。根据 API 的不同，运算可以在 CPU、支持 CUDA 的 GPU 或分布式多 GPU、多节点系统上运行。

nvmath-python v1.0 版本发布

随着 nvmath-python v1.0 的正式发布，本文将探讨该库的设计及其在加速数学运算方面的独特功能——从 CPU 或单 GPU 扩展到多 GPU、多节点规模。nvmath-python 是 CUDA 和 NVPL 数学库（如 cuFFT、cuBLASLt、cuDSS、cuSPARSE、cuTENSOR、cuBLASMp 等）的 Pythonic 抽象层。一种新颖的稀疏方法——通用稀疏张量（UST）——使用户能够通过特定领域语言创建自己独特的应用最优稀疏格式，而无需在代码中实现它。

快速灵活的安装

安装具有复杂原生依赖的 Python 包可能是一项耗时且令人沮丧的体验。nvmath-python 安装快速，且可针对不同环境进行定制。

选择包管理器，如 pip、conda、uv 或 pixi。
可以通过包管理器的依赖解析系统安装所有必需的依赖项，或执行最小安装，这在 CI/CD 或仅 CPU 环境等场景中非常有用。
选择 CPU 后端、设备 API 支持或分布式 API。
选择一个配套的数组库进行协作，如 NumPy、CuPy 或 PyTorch（或全部）。有关可用选项，请参阅详细的安装指南。

现有数组库的有益补充

与 NumPy 等其他数学库类似，nvmath-python 实现了许多工程和科学计算应用中常用的核心数值运算。然而，它并不旨在取代通用数组库或提供索引、切片或归约等传统功能。

相反，nvmath-python 专注于在 Python 中公开 CUDA-X 数学库的全部功能和强大性能，使现有数组库和框架能够更轻松地使用高度优化的 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)  # a is a NumPy array
b = np.random.randn(k, n)  # b is a NumPy array
c = nvmath.linalg.advanced.matmul(a, b)  # c is also a NumPy array

内存和执行空间的选择

选择数组库的灵活性同时适用于 GPU 库（如 CuPy）和 CPU 库（如 NumPy）。这是因为 nvmath-python 由以下内容提供支持：

GPU 库，如 cuBLAS 和 cuFFT。
CPU 库，如适用于 NVIDIA Grace 或任何 ARM v8 CPU 的 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 提供了一套全面的功能和配置，专门针对狭窄的操作范围设计，可能仅限于特定硬件。

举例来说，高级矩阵乘法针对 GPU 上的密集操作数专门实现了复合运算 \(\scriptstyle \mathbf{D}=f(\mathbf{A}\mathbf{B}+\mathbf{C})\)，并提供了挤压最高硬件效率所需的每一种配置。相反，通用矩阵乘法 API 支持跨 CPU 和 GPU 执行空间的密集和结构化操作数，但仅提供适用于其更广泛范围的通用选项子集。

最佳选择完全取决于具体用例：当某个操作成为计算瓶颈，需要硬件特定优化或访问特定功能时，专用 API 是理想之选。同时，对于性能要求不高或不需要专门定制的任务，通用 API 更合适。所有专用 API 都位于 advanced 子模块中，以将其与通用 API 区分开来。

使用 nvmath-python 进行日志记录

该库与 logging 模块的 Python 标准库记录器集成，用于捕获各个级别（调试、信息、警告和错误）的计算细节。

以下示例展示了内存和执行空间之间的数据流（使用高级 matmul）。

import numpy as np
import nvmath
import logging

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)

m, n, k = 8000, 2000, 4000
a_cpu = np.random.randn(m, k).astype(np.float32)
b_cpu = np.random.randn(k, n).astype(np.float32)
d_cpu = nvmath.linalg.advanced.matmul(a_cpu, b_cpu)

生成的输出将如下所示：

2025-09-18 14:53:32,166 INFO     = SPECIFICATION PHASE =
2025-09-18 14:53:32,167 INFO     The data type of operand A is 'float32', and that of operand B is 'float32'.
2025-09-18 14:53:32,168 INFO     The input operands' memory space is cpu, and the execution space is on device 0.
...

请注意显示操作数来源和消耗位置的记录。这是内存和执行空间之间可能存在昂贵数据传输的指示。现在运行一个使用通用 API（如 fft）的类似实验，以展示内存和执行空间之间的数据流。

import numpy as np
import nvmath
import logging

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)

N = 10000
e_cpu = (np.random.randn(N) + 1j * np.random.randn(N)).astype(np.complex64)
r_cpu = nvmath.fft.fft(e_cpu)

日志输出如下所示：

2025-09-18 15:46:22,295 INFO     The FFT type is C2C.
2025-09-18 15:46:22,295 INFO     The input data type is complex64, and the result data type is complex64.
2025-09-18 15:46:22,296 INFO     The specified FFT axes are (0,).
2025-09-18 15:46:22,297 INFO     The input tensor's memory space is cpu, and the execution space is cpu, with device cpu.
2025-09-18 15:46:22,298 INFO     The specified stream for the FFT ctor is None.
...

请注意，执行空间与输入的内存空间相同。只要有可能，nvmath-python 就会选择执行空间以最小化数据传输开销。用户可以通过向 API 提供 execution 关键字参数来自由选择所需的执行空间。

为什么复合运算很重要

像 \(\scriptstyle \mathbf{D}=f(\alpha\mathbf{A}\cdot\mathbf{B}+\beta\mathbf{C})\) 这样的运算，如果使用纯 NumPy 风格的 API，在许多用例中都能正常工作。然而，当底层基本运算的算术强度较低时，将它们串联成一系列调用是低效的。一个值得注意的例子是使用 \(\scriptstyle \mathbf{A}\) 作为高瘦矩阵计算 GEMM：

\(\scriptstyle \mathbf{D}=\alpha\mathbf{A}\cdot\mathbf{B}+\beta\mathbf{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 # Multiple kernels
d2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta) # Single kernel

图 1 显示，与 NumPy 风格的 API 相比，融合复合运算带来了可观的效益。

Significant performance speedup of fused kernels compared to unfused in low arithmetic intensity composite operations. — *图 1. 核融合对低算术强度复合运算的影响*

由于底层 cuBLASLt 库能够进行即时核融合，nvmath-python 的表现要好得多。这是提高算术强度的有效技术之一。

使用有状态 API 分摊准备成本

之前的所有示例都利用了 nvmath-python 的函数式或无状态 API。这是一个方便的单次调用 API，它涉及一个称为规划阶段的耗时准备逻辑。此外，准备成本还可能包括自动调优的成本。它与在规划/自动调优之后执行请求的数学运算的执行阶段不同。

性能说明

NVIDIA CUDA-X 数学库采用启发式方法来确定能产生最佳性能的特定实现。对于特定问题规模、布局或数据类型，可能有多种优化内核可供选择。哪个内核在特定硬件、工作负载和其他因素的组合上运行最佳并不总是显而易见的。自动调优旨在通过遍历内核选项、测量其性能并选择最佳选项来覆盖默认内核选择。因此，自动调优阶段可能非常耗时。

在深度学习等工作负载中，相同的操作可能会重复运行不同的输入。跨多次执行创建和重用计划可以分摊其规划成本。nvmath-python 的基于类的或有状态的 API 支持此工作流。

以下示例展示了在大小为 batch_size 的一批矩阵 a 和 b 以及偏置 bias 上使用 matmul 的类形式 API（带有 RELU_BIAS 尾部操作）。先前矩阵乘法的结果是下一个矩阵乘法的操作数，共有 feed_count 次操作。除了规划之外，它还执行自动调优阶段。

下面的代码展示了将有状态 API 的规划、自动调优和执行作为不同阶段使用。

import nvmath
from nvmath.linalg.advanced import MatmulEpilog
import cupy as cp

feed_count = 10  # The operation feed count.

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:
    # 1. Planning phase
    mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),
            epilog_inputs={"bias": bias})

    # 2. Autotuning phase
    mm.autotune(iterations=5)

    # 3. Execution phase.
    for i in range(feed_count):
        d = mm.execute()
        # The result of the previous MM is the operand `a` of the next MM, so use
        # reset_operands_unchecked() to reset the `a` operand.
        mm.reset_operands_unchecked(a=d)

Stateless APIs add a constant overhead to each operation. Stateful APIs quickly amortize one-time initial costs — *图 2. 有状态 API 通过多次执行快速分摊规划和自动调优成本*

图 2 显示了计算成本如何随执行次数变化。细虚线代表使用 nvmath-python 无状态 API 的成本。粗虚线显示切换到有状态 API 后的成本降低，点划线显示自动调优带来的额外性能增益。有状态 API 分摊了规范和准备成本，而无状态 API 在每次执行期间都会产生这些成本。自动调优的优势可以跨会话扩展，因为自动调优的计划可以序列化到磁盘并在新会话中加载。

图 3 显示，内置启发式方法通常可以在无需自动调优的情况下选择高性能内核。然而，某些问题规模、数据类型、操作数布局、硬件和其他因素的组合可以从自动调优中受益。在测试配置中，NVIDIA RTX A6000 显示出最大的加速，而 NVIDIA B200 无需自动调优即可达到峰值性能。

Bar chart showing the percentage speedup from autotuning on two GPUs: NVIDIA RTX A6000 shows over 250% speedup, while B200 shows minimal gains for this example (indicating that the heuristics do a great job). — *图 3. 执行阶段的自动调优并不总是能带来显著的加速，尽管 RTX A6000 在此特定问题上获得了 256% 的提升*

与 nvmath-python 融合的自定义内核

nvmath-python 与 numba-cuda 等 Python 编译器集成，支持将高性能自定义 Python 代码即时编译（JIT）并与 nvmath-python 运算一起使用。

自定义 FFT 回调

FFT 回调被编写为具有预定义签名的 Python 函数，并被 JIT 编译为中间表示，稍后用作 nvmath-python 正向或反向 FFT 的自定义前言或尾部操作。

Alt text: A sharp, grayscale, overhead photo of a light-colored dog sitting on pavement with its front paws elegantly crossed, casting a dark shadow. Then, the same photo of the dog now heavily blurred to demonstrate the effect of a Gaussian image filter. — *图 4. 左侧是应用高斯滤波器之前狗的原始灰度图像。右侧是使用 nvmath-python 的 FFT 和自定义 JIT 编译回调实现的结果图像*

高斯滤波器示例

作为说明，我们实现了一个高斯滤波器，它对原始图像应用模糊。下面的代码片段使用 PIL 库加载图像，然后将其转换为灰度 [0, 1] 图像作为 CuPy ndarray。对于图像滤波，我们实现了一个 img → R2C FFT → 高斯滤波器 → C2R iFFT → filtered_img 的链。高斯滤波器为 \(\scriptstyle G(x,y)=\exp\left(-\frac{x^2+y^2}{2\sigma^2}\right)\)，在频域中也是一个高斯函数 \(\scriptstyle H(f_x,f_y)=\exp\left(-2\pi^2\sigma^2(f_x^2+f_y^2)\right)\)。

以下代码展示了如何使用 nvmath-python FFT 和自定义回调函数应用高斯图像滤波器：

from PIL import Image
import nvmath
import cupy as cp

img = cp.asarray(Image.open("your_lovely_dog.jpg").convert("L")) / 255.0 # Gray[0,1]
wh = img.shape[0] * image.shape[1] # We must normalize by the image area
sigma_value = 20.0  # Filter size

# Implement Gaussian filter in the frequency domain
def gaussian_filter(shape, sigma):
    fy = cp.fft.fftfreq(shape[0])[:,None] # Column
    fx = cp.fft.rfftfreq(shape[1])[None,:] # Row
    return = cp.exp(-2.0 * cp.pi * cp.pi * sigma * sigma * (fx * fx + fy * fy))

# Implement FFT epilog wrapper with the pre-defined signature
def epilog_impl(data_out, offset, data, filter_data, unused): # Epilog to be compiled
    data_out[offset] = data * filter_data[offset] / wh

# Compile epilog to LTO-IR targeting the current CUDA device
epilog = nvmath.fft.compile_epilog(epilog_impl, "complex64", "complex64")

# Compute R2C FFT using nvmath-python with the compiled epilog
h_filter = gaussian_filter(img.shape, sigma)
img_fft = nvmath.fft.rfft(image, epilog={"ltoir": epilog, "data": h_filter.data.ptr})

# Compute C2R inverse FFT using nvmath-python
filtered_img = nvmath.fft.irfft(img_fft) # Visualize or save as you want

带有 nvmath-python 调用的自定义 numba-cuda 内核

第二个常用场景是在用 numba-cuda 编写的 GPU 内核中调用 nvmath-python 设备 API。nvmath-python 支持 FFT、GEMM、密集直接求解器（LU、Cholesky、QR）和 RNG 的设备 API。以下示例展示了用于蒙特卡洛股价模拟的几何布朗运动 (GBM) 的实现。它使用 nvmath-python 的随机数生成器生成高斯分布，以及将正态分布转换为 GBM 蒙特卡洛路径的自定义 numba-cuda 代码：

from numba import cuda
from nvmath.device import random
import cupy as cp
import math

# Pre-compile the RNGs into IR to use alongside other device code
compiled_rng = random.Compile(cc=None)

# GBM parameters
rng_seed = 7777
n_time_steps, n_paths = 252, 8192
mu, sigma, s0 = 0.003, 0.027, 100.0

# Set up CUDA kernel launch configuration
threads_per_block = 32
blocks = n_paths // threads_per_block + bool(n_paths % threads_per_block)
nthreads = threads_per_block * blocks

# RNG initialization kernel
@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])

# GBM path generation kernel
@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

    # Consume 4 normal variates at a time for better throughput
    for i in range(1, nsteps, 4):
        v = random.normal4(states[idx])  # Returned as float32x4 type
        vals = v.x, v.y, v.z, v.w  # Decompose into a tuple of float32
        for j in range(i, min(i + 4, nsteps)):  # Process a chunk of 4 time steps
            paths[idx, j] = paths[idx, j - 1] * math.exp(mu + sigma * vals[j - i])

# Initialize RNG
states = random.StatesPhilox4_32_10(nthreads)
init_rng[blocks, threads_per_block](states, rng_seed)

# Generate GBM paths on GPU
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)

generate_gbm_paths 中的每个操作都具有低算术强度，这使得基于主机 API 的实现效率低下。将这些操作与 numba-cuda 和 nvmath-python 设备 API 融合至关重要。

开始使用 nvmath-python

nvmath-python 的设计旨在提高生产力，同时不牺牲性能，重新构想了现代数学库的设计。只需一个简单的命令即可开始使用：

pip install nvmath-python[cu13]

其他资源包括：

包含详细设置说明的安装文档
nvmath-python GitHub 仓库，包括代码示例和深入的教程笔记本
NVIDIA 加速计算中心中的扩展 Python 培训材料
之前的文章，使用 nvmath-python 中的通用稀疏张量简化稀疏深度学习
nvmath-python 产品页面，了解最新消息和公告

致谢

该库是 NVIDIA 众多人员共同努力的结果，包括：

Harun Bayraktar, Becca Zandstein, Lukasz Ligowski, Aart Bik, Yevhenii Havrylko, Juan Galvez, Daniel Ching, Mark Olah, Yang Gao, Szymon Karpinski , Kamil Tokarski , Francesco Rizzi, Jakub Lisowski, Marcin Rogowski, Robbie Jensen , Artem Amogolonov, Sushma Kini, Rachna Pandey, Graham Markall, Michael Yh Wang, Bradley Dice, Liam Zhang, Jack Cui, Chang Liu, Qi Xia, Feng Cheng, Ruilin Tian, Zan Xu, Almog Segal, Kirill Voronin, Evarist Fomenko, 以及更多贡献者。

Read the original source Next news

#Performance