NVIDIA nvmath-python is a library designed to bridge the gap between the Python scientific community and NVIDIA CUDA-X math libraries. It gives Python users access to CUDA-X performance for common math operations without disrupting existing workflows. Depending on the API, operations can run on a CPU, CUDA-enabled GPU, or distributed multi-GPU, multi-node systems.
nvmath-python v1.0 release
With the general availability of nvmath-python v1.0, this post explores the library’s design and unique capabilities for accelerating math operations—from a CPU or single GPU up to multi-GPU, multi-node scale. nvmath-python is a Pythonic abstraction layer over the CUDA and NVPL math libraries such as cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR, cuBLASMp, and more. A novel approach to sparsity, the universal sparse tensor (UST), enables the user to create their own unique application-optimal sparse format through a domain-specific language without having to implement it in code.
Fast and flexible installation
Installing a Python package with complex native dependencies can be a time-consuming and frustrating experience. nvmath-python installs quickly and can be customized for different environments.
- Choose a package manager, such as pip, conda, uv, or pixi.
- There is an option to install all required dependencies through the package manager’s dependency resolution system or perform a bare-minimum installation, useful in scenarios such as CI/CD or CPU-only environments.
- Pick and choose the CPU backend, device APIs support, or distributed APIs.
- Choose a companion array library to work with, such as NumPy, CuPy, or PyTorch (or all of them). See the detailed installation guide for available options.
A useful complement to existing array libraries
Like other math libraries such as NumPy, nvmath-python implements core numerical operations useful in many engineering and scientific computing applications. However, it’s not intended to replace general-purpose array libraries or provide traditional features like indexing, slicing, or reduction.
Instead, nvmath-python focuses on exposing the full functionality and power of CUDA-X math libraries in Python, making it easier for existing array libraries and frameworks to use highly optimized GPU-accelerated routines without relying on low-level C/C++ interfaces.
In the following example, nvmath-python consumes NumPy arrays and the result is also a NumPy array.
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
Choice of memory and execution spaces
The flexibility of choosing an array library applies to both GPU libraries, such as CuPy, and CPU libraries, such as NumPy. This is possible because nvmath-python is backed by the following:
- GPU libraries such as cuBLAS and cuFFT.
- CPU libraries such as NVPL for NVIDIA Grace or any ARM v8 CPUs and Intel MKL for x86 hosts.
- Distributed libraries such as cuBLASMp, cuSOLVERMp, or cuFFTMp.
This support simplifies code migration between CPU and GPU and enables hybrid and distributed workflows that combine CPU and GPU execution.
The following code illustrates how nvmath-python supports multiple memory and execution spaces.
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 fft execution space for each call is inferred from its input tensor, either a_gpu or a_cpu, although a different execution space can be specified. The library’s logging facility shows where each operation ran.
Generic and specialized APIs
The APIs within nvmath-python are broadly divided into two classes: generic APIs that act as flexible multitools (wide but shallow), and specialized APIs designed as precise, dedicated instruments (narrow and deep).
Generic APIs focus on providing a uniform user experience across various execution and memory spaces as well as operand types, however they restrict configurability to the baseline, common features shared across their broad scope. Meanwhile, specialized APIs provide a comprehensive set of features and configurations designed specifically for a narrow operational range and may be restricted to particular hardware.
To illustrate, the advanced matrix multiplication implements the composite operation \(\scriptstyle \mathbf{D}=f(\mathbf{A}\mathbf{B}+\mathbf{C})\) specifically for dense operands on the GPU and provides every configuration necessary to squeeze out the highest possible hardware efficiency. Conversely, the generic matrix multiplication API accommodates dense and structured operands across CPU and GPU execution spaces, but offers only the common subset of options applicable to its wider scope.
The optimal choice depends entirely on the specific use-case: specialized APIs are ideal when an operation becomes a computational bottleneck that demands hardware-specific optimizations or access to distinct features. Meanwhile, generic APIs are better suited for tasks that are not performance-critical or when specialized customization is unnecessary. All specialized APIs live within the advanced submodules to keep them distinct from generic APIs.
Logging with nvmath-python
The library provides integration with the Python standard library logger from the logging module for capturing computational details at various levels (debug, information, warning, and error).
The following example illustrates the data flow between memory and execution spaces (using the advanced 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)
The produced output will look like:
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.
...
Take note of the record showing where operands come from and where they are consumed. This is an indication of potentially expensive data transfer between memory and execution spaces. Now run a similar experiment with a generic API like fft to illustrate data flow between memory and execution spaces.
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)
The logging output looks like:
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.
...
Note that execution space is the same as inputs’ memory space. Whenever possible, nvmath-python selects the execution space to minimize the data transfer overheads. The user is free to select the desired execution space by providing the execution keyword argument to an API.
Why composite operations matter
An operation like \(\scriptstyle \mathbf{D}=f(\alpha\mathbf{A}\cdot\mathbf{B}+\beta\mathbf{C})\) with a pure NumPy-like API will work decently in many use cases. However, when underlying primitive operations have low arithmetic intensity, chaining them as a series of calls is inefficient. A notable example is computing GEMM with \(\scriptstyle \mathbf{A}\) being a tall-and-skinny matrix:
\(\scriptstyle \mathbf{D}=\alpha\mathbf{A}\cdot\mathbf{B}+\beta\mathbf{C}\)
The following code illustrates GEMM on tall-and-skinny matrices with CuPy and nvmath-python.
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
Figure 1 shows that a fused composite operation brings measurable benefits compared to NumPy-like APIs.
nvmath-python performs much better due to the underlying cuBLASLt library, capable of just-in-time kernel fusion. It is among the effective techniques for increasing arithmetic intensity.
Amortizing preparation costs by using stateful APIs
All previous examples exploit the functional-form, or stateless, API of the nvmath-python. It’s a convenient single-call API, which involves a time-consuming preparation logic, called the planning phase. Additionally the preparation cost may also include the cost of autotuning. It is distinct from the execution phase that performs the requested math operation after the planning/autotuning.
Performance note
NVIDIA CUDA-X math libraries employ heuristics to determine a specific implementation that yields the best performance. There can be multiple choices of specialized kernels optimized for certain problem sizes, layouts or data types. It is not always obvious which kernel will run best on a specific combination of hardware, workload and other factors. Autotuning aims at overriding the default kernel selection by iterating through kernel options, measuring their performance, and choosing the best one. As a result, the autotuning phase may be very time consuming.
In workloads such as deep learning, the same operation may run repeatedly with different inputs. Creating and reusing a plan across executions amortizes its planning cost. nvmath-python’s class-based, or stateful, APIs support this workflow.
The following example illustrates the use of class-form API for matmul (with RELU_BIAS epilog) on a batch of the batch_size size of matrices a and b, and biases bias. The result of the prior matrix multiplication is an operand in the next matrix multiplication, and there are feed_count operations. Besides planning it also performs an autotuning phase.
The code below illustrates using stateful APIs with planning, autotuning, and execution as distinct phases.
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)
Figure 2 shows how computational cost changes with the number of executions. The fine dashed line represents the cost of using nvmath-python’s stateless API. The coarse dashed line shows the cost reduction from switching to the stateful API, and the dash-dot line shows the additional performance gain from autotuning. The stateful API amortizes specification and preparation costs, while the stateless API incurs them during every execution. Autotuning benefits can extend across sessions because an autotuned plan can be serialized to disk and loaded in a new session.
Figure 3 shows that built-in heuristics can often select a high-performing kernel without autotuning. However, some combinations of problem size, data type, operand layout, hardware, and other factors benefit from autotuning. In the tested configuration, the NVIDIA RTX A6000 shows the largest speedup, while the NVIDIA B200 reaches peak performance without autotuning.
Deployment note
Another example is performing the tuning once and then deploying the tuned plan across multiple homogeneous systems to perform a similar operation on data that doesn’t require repetitive replanning. For a deeper dive please refer to example14_autotune.py, example15_manual_tuning.py, and example16_reuse_algorithms.py in the nvmath-python GitHub repo.
Custom kernels fused with nvmath-python
nvmath-python integrates with Python compilers such as numba-cuda, enabling high-performance custom Python code to be compiled just in time (JIT) and used alongside nvmath-python operations.
Custom FFT callbacks
Callbacks for FFT are written as Python functions with a predefined signature and JIT-compiled to intermediate representation, which is later used as a custom prolog or epilog for nvmath-python’s forward or inverse FFT.
Gaussian filter example
As an illustration we implement a Gaussian filter, which applies blurring to the original image. The below code snippet uses PIL library for image loading, which is then converted to a grayscale [0, 1] image as a CuPy ndarray. For image filtration we implement a chain of img → R2C FFT → Gaussian filter → C2R iFFT → filtered_img. The Gaussian filter is \(\scriptstyle G(x,y)=\exp\left(-\frac{x^2+y^2}{2\sigma^2}\right)\), which in frequency domain is also a Gaussian \(\scriptstyle H(f_x,f_y)=\exp\left(-2\pi^2\sigma^2(f_x^2+f_y^2)\right)\).
The following code shows how to apply a Gaussian image filter with nvmath-python FFT and custom callback function:
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
Custom numba-cuda kernels with nvmath-python calls
The second commonly used scenario is calling nvmath-python device APIs from within GPU kernels written in numba-cuda. nvmath-python supports device APIs for FFTs, GEMM, dense direct solvers (LU, Cholesky, QR) and RNG. The following example shows the implementation of the Geometric Brownian Motion (GBM) for Monte Carlo stock price simulations. It uses nvmath-python’s random number generator for Gaussian distribution along with the custom numba-cuda code that converts normal distribution to the GBM Monte Carlo paths:
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)
Every operation in generate_gbm_paths has low arithmetic intensity, which makes the host API-based implementation inefficient. It is crucial to get these operations fused with numba-cuda and nvmath-python device APIs.
Get started with nvmath-python
Designed for productivity without performance compromises, nvmath-python reimagines the design of modern math libraries. Get started with one simple command:
pip install nvmath-python[cu13]
Additional resources include:
- Installation documentation with detailed setup instructions
- The nvmath-python GitHub repository, including code examples and in-depth tutorial notebooks
- Extended Python training materials in the NVIDIA Accelerated Computing Hub
- The earlier post, Simplify Sparse Deep Learning with the Universal Sparse Tensor in nvmath-python
- The nvmath-python product page for the latest news and announcements
Acknowledgments
The library is a result of efforts of many people from across NVIDIA, 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 many more.
