NVIDIA、Python科学コミュニティ向け数学ライブラリ「nvmath-python」1.0 を一般公開

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)

b = np.random.randn(k, n)

c = nvmath.linalg.advanced.matmul(a, b)

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 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 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 being a tall-and-skinny matrix:

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

d2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta)

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

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)

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.

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 , which in frequency domain is also a Gaussian .

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

wh = img.shape[0] * image.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)

img_fft = nvmath.fft.rfft(image, epilog={"ltoir": epilog, "data": h_filter.data.ptr})

filtered_img = nvmath.fft.irfft(img_fft)

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

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)

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.*