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

NVIDIA released nvmath-python v1.0, a Python library that provides access to NVIDIA CUDA-X math libraries for high-performance core math operations on CPU, single GPU, or distributed multi-GPU, multi-node systems. The library supports flexible installation via pip, conda, uv, or pixi, and works with array libraries like NumPy, CuPy, and PyTorch, enabling seamless code migration between CPU and GPU execution.

NVIDIA nvmath-python https://developer.nvidia.com/nvmath-python is a library designed to bridge the gap between the Python scientific community and NVIDIA CUDA-X math libraries https://developer.nvidia.com/gpu-accelerated-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 https://developer.nvidia.com/blog/simplify-sparse-deep-learning-with-universal-sparse-tensor-in-nvmath-python/ 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 https://pypackaging-native.github.io . 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 https://docs.nvidia.com/cuda/nvmath-python/latest/installation.html 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. python 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 https://developer.nvidia.com/cublas and cuFFT https://developer.nvidia.com/cufft . - CPU libraries such as NVPL https://developer.nvidia.com/nvpl for NVIDIA Grace or any ARM v8 CPUs and Intel MKL for x86 hosts. - Distributed libraries such as cuBLASMp https://docs.nvidia.com/cuda/cublasmp/index.html , cuSOLVERMp https://docs.nvidia.com/cuda/cusolvermp/index.html , or cuFFTMp https://docs.nvidia.com/cuda/cufftmp/index.html . 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. 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 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 https://github.com/NVIDIA/nvmath-python/blob/main/examples/fft/example01 numpy.py . The library’s logging facility https://docs.nvidia.com/cuda/nvmath-python/latest/overview.html full-logging-support 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 https://docs.python.org/3/library/logging.html 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 . python 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. python 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 https://arxiv.org/abs/2002.03258 : \ \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. 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 https://docs.nvidia.com/cuda/cublas/ using-the-cublaslt-api 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. python 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 https://github.com/NVIDIA/nvmath-python/blob/main/examples/linalg/advanced/matmul/example14 autotune.py , , and https://github.com/NVIDIA/nvmath-python/blob/main/examples/linalg/advanced/matmul/example15 manual tuning.py example15 manual tuning.py in the nvmath-python GitHub repo. https://github.com/NVIDIA/nvmath-python/blob/main/examples/linalg/advanced/matmul/example16 reuse algorithms.py example16 reuse algorithms.py 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 C allbacks 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: python 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 https://en.wikipedia.org/wiki/Geometric Brownian motion 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: python 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 https://docs.nvidia.com/cuda/nvmath-python/latest/installation.html with detailed setup instructions- The nvmath-python GitHub repository https://github.com/NVIDIA/nvmath-python , including code examples https://github.com/NVIDIA/nvmath-python/tree/main/examples and in-depth tutorial notebooks https://github.com/NVIDIA/nvmath-python/tree/main/notebooks Extended Python training materials https://github.com/NVIDIA/accelerated-computing-hub/tree/main in the NVIDIA Accelerated Computing Hub- The earlier post, Simplify Sparse Deep Learning with the Universal Sparse Tensor in nvmath-python https://developer.nvidia.com/blog/simplify-sparse-deep-learning-with-universal-sparse-tensor-in-nvmath-python/ - The nvmath-python product page https://developer.nvidia.com/nvmath-python 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.