NVIDIA has made nvmath-python v1.0 generally available: a Pythonic abstraction layer that exposes CUDA‑X math libraries (for example cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR, cuBLASMp) to Python developers. The library is designed so users can access high‑performance CUDA‑X routines from Python without rewriting code in C/C++. Depending on the API used, operations can execute on a CPU, a single CUDA‑enabled GPU, or distributed multi‑GPU/multi‑node systems.
What the library provides
nvmath-python aims to provide two complementary capabilities. First, it implements core numerical operations commonly needed in engineering and scientific computing. Second, it exposes the full functionality of CUDA‑X math libraries to Python users, enabling existing array libraries and frameworks to call highly optimized GPU and distributed routines without resorting to low‑level C/C++ interfaces. The project also introduces a universal sparse tensor (UST) approach that lets users define application‑optimal sparse formats via a domain‑specific language rather than implementing them manually in low‑level code.
Fast, flexible installation
The package is intended to simplify installing Python packages that have complex native dependencies. Installers can use pip, conda, uv, or pixi and choose between a full dependency resolution installation or a bare‑minimum install suitable for CI/CD or CPU‑only environments. Users can pick CPU backends (for example NVPL for NVIDIA Grace or Intel MKL for x86), device API support, distributed APIs, and companion array libraries such as NumPy, CuPy, or PyTorch. A quick install example shown by NVIDIA is:
pip install nvmath-python[cu13]
Works alongside existing array libraries
nvmath-python is intended to complement, not replace, general‑purpose array libraries. Inputs and outputs can be NumPy, CuPy or PyTorch arrays (or combinations). This makes migrating code between CPU and GPU straightforward and supports hybrid or distributed workflows.
Example: a NumPy input returns a NumPy result:
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)
Memory and execution spaces
nvmath-python can infer and manage both memory and execution spaces. GPU libraries such as cuBLAS and cuFFT, CPU backends like NVPL or Intel MKL, and distributed libraries such as cuBLASMp, cuSOLVERMp, or cuFFTMp are supported. The library attempts to select an execution space that minimizes data transfer overheads; callers can also explicitly specify the execution space.
Example showing both CPU and GPU FFT calls:
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)
The library’s logging facility reports where each operation runs and where operands live.
Generic vs. specialized APIs
APIs are grouped into generic (broad scope, limited configurability) and specialized (narrow scope, deep configurability) classes. Generic APIs provide a uniform interface across execution and memory spaces and support a variety of operand types but expose only common options. Specialized APIs (contained in the advanced submodules) expose full hardware‑specific controls and are intended for performance‑critical operations where fine tuning matters.
For instance, a specialized advanced matmul implements D = f(A B + C) for dense GPU operands with every configuration option to maximize hardware efficiency, whereas the generic matmul API supports dense and structured operands across CPU and GPU but only exposes the common subset of options.
Logging and visibility into data flow
nvmath-python integrates with Python’s standard logging module and can emit debug/info/warning/error messages that describe the specification and execution phases. Example logs in the documentation (timestamped 2025‑09‑18 in the examples) show messages such as operand data types, memory spaces, and chosen execution devices—information useful to identify expensive data transfers between memory and execution spaces.
Where possible, the library chooses execution spaces to avoid unnecessary copies, but users can override selection via an execution keyword argument.
Why fused composite operations help
Chaining low‑arithmetic‑intensity primitives (like separate multiplies and adds) can be inefficient. nvmath-python leverages cuBLASLt’s just‑in‑time kernel fusion to execute composite operations in a single kernel, increasing arithmetic intensity and improving performance, particularly for cases like tall‑and‑skinny GEMM.
Example: for m = 10_000_000, k = 10, n = 40, a fused advanced.matmul call executes as a single kernel, while the NumPy/CuPy style expression performs multiple kernels and incurs extra overhead.
Stateful APIs, planning and autotuning
Stateless (functional‑form) APIs perform planning and optionally autotuning as part of each call. Planning and autotuning can be time‑consuming, so for repeated workloads the library provides class‑based stateful APIs where planning, autotuning and execution are distinct phases. A plan can be autotuned once and reused many times to amortize the preparation cost; autotuned plans can also be serialized to disk and loaded in another session.
Documentation examples show: create a Matmul object, call plan(...) to specify epilog/inputs, run autotune(iterations=...), then call execute() multiple times. Performance graphs in the documentation illustrate how stateful APIs reduce per‑execution cost, while autotuning can further improve runtime for some GPU/hardware/problem combinations.
NVIDIA notes that autotuning is not always necessary; built‑in heuristics often choose a high‑performing kernel. In the provided measurements, the NVIDIA RTX A6000 showed a significant speedup with autotuning for a particular problem (reported as a 256% increase in that test), whereas another GPU (NVIDIA B200 in the examples) reached peak performance without autotuning.
Custom kernels and JIT callbacks
nvmath-python supports custom JIT‑compiled callbacks (for example for FFT epilog/prolog) and integrates with compilers such as numba‑cuda so that user‑defined device code and nvmath device APIs can be combined. Examples in the documentation include:
-
A Gaussian image filter implemented by applying R2C FFT → multiply by Gaussian filter in frequency domain → C2R iFFT using an nvmath‑compiled FFT epilog. The epilog is written as a Python function, compiled to intermediate representation (LTO‑IR) and passed to the FFT call.
-
Device code examples where numba‑cuda kernels call nvmath device RNGs. One example implements Geometric Brownian Motion (GBM) Monte Carlo paths: RNG state initialization and path generation kernels are JIT‑compiled and linked with compiled RNG IR so that RNGs run efficiently on the GPU.
These integrations are useful when per‑operation arithmetic intensity is low and fusing logic into device code matters for throughput.
Deployment and further resources
NVIDIA recommends autotuning once and reusing tuned plans across homogeneous systems for repetitive workloads. The repository contains practical examples (example14_autotune.py, example15_manual_tuning.py, example16_reuse_algorithms.py) to illustrate tuning and reuse patterns.
Getting started: pip install nvmath‑python[cu13]. Additional materials include an installation guide, the nvmath‑python GitHub repository with examples and tutorial notebooks, extended Python training on the NVIDIA Accelerated Computing Hub, and a previous post on the universal sparse tensor.
Acknowledgments
NVIDIA credits many contributors to the library, including 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 and others.
