NVIDIA nvmath-python is a library designed to bridge the hole between the Python scientific group and NVIDIA CUDA-X math libraries. It offers Python customers entry to CUDA-X efficiency for widespread math operations with out disrupting current workflows. Relying on the API, operations can run on a CPU, CUDA-enabled GPU, or distributed multi-GPU, multi-node programs.<\/p>\n
nvmath-python v1.0 launch<\/h2>\n
With the overall availability of nvmath-python v1.0, this publish explores the library\u2019s design and distinctive capabilities for accelerating math operations\u2014from a CPU or single GPU as much as multi-GPU, multi-node scale. nvmath-python is a Pythonic abstraction layer over the CUDA and NVPL math libraries resembling cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR, cuBLASMp, and extra. A novel method to sparsity, the common sparse tensor (UST), permits the person to create their very own distinctive application-optimal sparse format via a domain-specific language with out having to implement it in code.<\/p>\n
Quick and versatile set up<\/h2>\n
Putting in a Python bundle with complicated native dependencies is usually a time-consuming and irritating expertise. nvmath-python installs shortly and might be custom-made for various environments.<\/p>\n
Select a bundle supervisor, resembling pip, conda, uv, or pixi.<\/p>\n
There may be an choice to put in all required dependencies via the bundle supervisor\u2019s dependency decision system or carry out a bare-minimum set up, helpful in situations resembling CI\/CD or CPU-only environments.<\/p>\n
Decide and select the CPU backend, system APIs help, or distributed APIs.<\/p>\n
Select a companion array library to work with, resembling NumPy, CuPy, or PyTorch (or all of them). See the detailed set up information for out there choices.<\/p>\n
A helpful complement to current array libraries<\/h2>\n
Like different math libraries resembling NumPy, nvmath-python implements core numerical operations helpful in lots of engineering and scientific computing purposes. Nonetheless, it\u2019s not supposed to switch general-purpose array libraries or present conventional options like indexing, slicing, or discount.<\/p>\n
As a substitute, nvmath-python focuses on exposing the complete performance and energy of CUDA-X math libraries in Python, making it simpler for current array libraries and frameworks to make use of extremely optimized GPU-accelerated routines with out counting on low-level C\/C++ interfaces.<\/p>\n
Within the following instance, nvmath-python consumes NumPy arrays and the consequence can be a NumPy array.<\/p>\n
\nimport nvmath<\/p>\n
m, n, ok = 10, 40, 100
\na = np.random.randn(m, ok) # a is a NumPy array
\nb = np.random.randn(ok, n) # b is a NumPy array
\nc = nvmath.linalg.superior.matmul(a, b) # c can be a NumPy array\n<\/p><\/div>\n
Selection of reminiscence and execution areas<\/h2>\n
The flexibleness of selecting an array library applies to each GPU libraries, resembling CuPy, and CPU libraries, resembling NumPy. That is attainable as a result of nvmath-python is backed by the next:<\/p>\n
This help simplifies code migration between CPU and GPU and permits hybrid and distributed workflows that mix CPU and GPU execution.<\/p>\n
The next code illustrates how nvmath-python helps a number of reminiscence and execution areas.<\/p>\n
\nimport numpy as np
\nimport nvmath<\/p>\n
N = 2048
\na_gpu = cp.random.randn(N) + 1j * cp.random.randn(N)
\na_cpu = np.random.randn(N) + 1j * np.random.randn(N)
\nc_gpu = nvmath.fft.fft(a_gpu)
\nc_cpu = nvmath.fft.fft(a_cpu)\n<\/p><\/div>\n
The fft execution house for every name is inferred from its enter tensor, both a_gpu or a_cpu, though a special execution house might be specified. The library\u2019s logging facility exhibits the place every operation ran.<\/p>\n
Generic and specialised APIs<\/h2>\n
The APIs inside nvmath-python are broadly divided into two lessons: generic APIs that act as versatile multitools (vast however shallow), and specialised APIs designed as exact, devoted devices (slim and deep).<\/p>\n
Generic APIs give attention to offering a uniform person expertise throughout numerous execution and reminiscence areas in addition to operand sorts, nevertheless they prohibit configurability to the baseline, widespread options shared throughout their broad scope. In the meantime, specialised APIs present a complete set of options and configurations designed particularly for a slim operational vary and could also be restricted to explicit {hardware}. <\/p>\n
As an instance, the superior matrix multiplication implements the composite operation (scriptstyle mathbf{D}=f(mathbf{A}mathbf{B}+mathbf{C})) particularly for dense operands on the GPU and gives each configuration essential to squeeze out the very best attainable {hardware} effectivity. Conversely, the generic matrix multiplication API accommodates dense and structured operands throughout CPU and GPU execution areas, however provides solely the widespread subset of choices relevant to its wider scope.<\/p>\n
The optimum alternative relies upon completely on the precise use-case: specialised APIs are ultimate when an operation turns into a computational bottleneck that calls for hardware-specific optimizations or entry to distinct options. In the meantime, generic APIs are higher fitted to duties that aren’t performance-critical or when specialised customization is pointless. All specialised APIs stay throughout the superior submodules to maintain them distinct from generic APIs.<\/p>\n
Logging with nvmath-python<\/h2>\n
The library gives integration with the Python customary library logger from the logging module for capturing computational particulars at numerous ranges (debug, info, warning, and error).<\/p>\n
The next instance illustrates the info movement between reminiscence and execution areas (utilizing the superior matmul).<\/p>\n
\nimport nvmath
\nimport logging<\/p>\n
logging.basicConfig(degree=logging.INFO,
\n format=”%(asctime)s %(levelname)-8s %(message)s”, pressure=True)
\nlogging.disable(logging.NOTSET)<\/p>\n
m, n, ok = 8000, 2000, 4000
\na_cpu = np.random.randn(m, ok).astype(np.float32)
\nb_cpu = np.random.randn(ok, n).astype(np.float32)
\nd_cpu = nvmath.linalg.superior.matmul(a_cpu, b_cpu)\n<\/p><\/div>\n
The produced output will seem like:<\/p>\n
\n2025-09-18 14:53:32,167 INFO The information sort of operand A is ‘float32’, and that of operand B is ‘float32′.
\n2025-09-18 14:53:32,168 INFO The enter operands’ reminiscence house is cpu, and the execution house is on system 0.
\n…\n<\/div>\n
Be aware of the document exhibiting the place operands come from and the place they’re consumed. This is a sign of doubtless costly knowledge switch between reminiscence and execution areas. Now run an analogous experiment with a generic API like fft for instance knowledge movement between reminiscence and execution areas.<\/p>\n
\nimport nvmath
\nimport logging<\/p>\n
logging.basicConfig(degree=logging.INFO,
\n format=”%(asctime)s %(levelname)-8s %(message)s”, pressure=True)
\nlogging.disable(logging.NOTSET)<\/p>\n
N = 10000
\ne_cpu = (np.random.randn(N) + 1j * np.random.randn(N)).astype(np.complex64)
\nr_cpu = nvmath.fft.fft(e_cpu)\n<\/p><\/div>\n
The logging output appears to be like like:<\/p>\n
\n2025-09-18 15:46:22,295 INFO The enter knowledge sort is complex64, and the consequence knowledge sort is complex64.
\n2025-09-18 15:46:22,296 INFO The desired FFT axes are (0,).
\n2025-09-18 15:46:22,297 INFO The enter tensor’s reminiscence house is cpu, and the execution house is cpu, with system cpu.
\n2025-09-18 15:46:22,298 INFO The desired stream for the FFT ctor is None.
\n…\n<\/div>\n
Word that execution house is identical as inputs\u2019 reminiscence house. Every time attainable, nvmath-python selects the execution house to reduce the info switch overheads. The person is free to pick the specified execution house by offering the execution key phrase argument to an API.<\/p>\n
Why composite operations matter<\/h2>\n
An operation like (scriptstyle mathbf{D}=f(alphamathbf{A}cdotmathbf{B}+betamathbf{C})) with a pure NumPy-like API will work decently in lots of use circumstances. Nonetheless, when underlying primitive operations have low arithmetic depth, chaining them as a collection of calls is inefficient. A notable instance is computing GEMM with (scriptstyle mathbf{A}) being a tall-and-skinny matrix:<\/p>\n
(scriptstyle mathbf{D}=alphamathbf{A}cdotmathbf{B}+betamathbf{C})<\/p>\n
The next code illustrates GEMM on tall-and-skinny matrices with CuPy and nvmath-python.<\/p>\n
\nimport nvmath<\/p>\n
m, n, ok = 10_000_000, 40, 10
\na = cp.random.randn(m, ok, dtype=cp.float32)
\nb = cp.random.randn(ok, n, dtype=cp.float32)
\nc = cp.random.randn(m, n, dtype=cp.float32)
\nalpha, beta = 1.5, 0.5<\/p>\n
d1 = alpha * cp.matmul(a, b) + beta * c # A number of kernels
\nd2 = nvmath.linalg.superior.matmul(a, b, c=c, alpha=alpha, beta=beta) # Single kernel\n<\/p><\/div>\n
Determine 1 exhibits {that a} fused composite operation brings measurable advantages in comparison with NumPy-like APIs.<\/p>\n
![\"Significant](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/image-81.webp\")
nvmath-python performs significantly better because of the underlying cuBLASLt library, able to just-in-time kernel fusion. It’s among the many efficient strategies for rising arithmetic depth.<\/p>\n
Amortizing preparation prices through the use of stateful APIs<\/h2>\n
All earlier examples exploit the functional-form, or stateless, API of the nvmath-python. It\u2019s a handy single-call API, which includes a time-consuming preparation logic, known as the planning part. Moreover the preparation price might also embrace the price of autotuning. It’s distinct from the execution part that performs the requested math operation after the planning\/autotuning.<\/p>\n
Efficiency word<\/p>\n
NVIDIA CUDA-X math libraries make use of heuristics to find out a particular implementation that yields the perfect efficiency. There might be a number of selections of specialised kernels optimized for sure drawback sizes, layouts or knowledge sorts. It’s not at all times apparent which kernel will run greatest on a particular mixture of {hardware}, workload and different elements. Autotuning goals at overriding the default kernel choice by iterating via kernel choices, measuring their efficiency, and selecting the perfect one. In consequence, the autotuning part could also be very time consuming.<\/p>\n<\/div>\n<\/div>\n
In workloads resembling deep studying, the identical operation could run repeatedly with completely different inputs. Creating and reusing a plan throughout executions amortizes its planning price. nvmath-python\u2019s class-based, or stateful, APIs help this workflow.<\/p>\n
The next instance illustrates using class-form API for matmul (with RELU_BIAS epilog) on a batch of the batch_size dimension of matrices a and b, and biases bias. The results of the prior matrix multiplication is an operand within the subsequent matrix multiplication, and there are feed_count operations. In addition to planning it additionally performs an autotuning part.<\/p>\n
The code beneath illustrates utilizing stateful APIs with planning, autotuning, and execution as distinct phases.<\/p>\n
\nfrom nvmath.linalg.superior import MatmulEpilog
\nimport cupy as cp<\/p>\n
feed_count = 10 # The operation feed depend.<\/p>\n
batch_size = 1024
\nm, n, ok = 1024, 1024, 1024<\/p>\n
a = cp.random.rand(batch_size, m, ok, dtype=cp.float32)
\nb = cp.random.rand(batch_size, ok, n, dtype=cp.float32)
\nbias = cp.random.rand(batch_size, m, 1, dtype=cp.float32)<\/p>\n
with nvmath.linalg.superior.Matmul(a, b) as mm:
\n # 1. Planning part
\n mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),
\n epilog_inputs={“bias”: bias})<\/p>\n
# 2. Autotuning part
\n mm.autotune(iterations=5)<\/p>\n
# 3. Execution part.
\n for i in vary(feed_count):
\n d = mm.execute()
\n # The results of the earlier MM is the operand `a` of the subsequent MM, so use
\n # reset_operands_unchecked() to reset the `a` operand.
\n mm.reset_operands_unchecked(a=d)\n<\/p><\/div>\n
![\"Stateless](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/image-82.webp\")
Determine 2 exhibits how computational price modifications with the variety of executions. The high-quality dashed line represents the price of utilizing nvmath-python\u2019s stateless API. The coarse dashed line exhibits the associated fee discount from switching to the stateful API, and the dash-dot line exhibits the extra efficiency achieve from autotuning. The stateful API amortizes specification and preparation prices, whereas the stateless API incurs them throughout each execution. Autotuning advantages can lengthen throughout periods as a result of an autotuned plan might be serialized to disk and loaded in a brand new session.<\/p>\n
Determine 3 exhibits that inbuilt heuristics can usually choose a high-performing kernel with out autotuning. Nonetheless, some mixtures of drawback dimension, knowledge sort, operand structure, {hardware}, and different elements profit from autotuning. Within the examined configuration, the NVIDIA RTX A6000 exhibits the biggest speedup, whereas the NVIDIA B200 reaches peak efficiency with out autotuning.<\/p>\n
![\"Bar](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/image-83.webp\")
Customized kernels fused with nvmath-python<\/h2>\n
nvmath-python integrates with Python compilers resembling numba-cuda, enabling high-performance customized Python code to be compiled simply in time (JIT) and used alongside nvmath-python operations.<\/p>\n
Customized FFT callbacks<\/h3>\n
Callbacks for FFT are written as Python capabilities with a predefined signature and JIT-compiled to intermediate illustration, which is later used as a customized prolog or epilog for nvmath-python\u2019s ahead or inverse FFT.<\/p>\n Gaussian filter instance<\/p>\n As an illustration we implement a Gaussian filter, which applies blurring to the unique picture. The beneath code snippet makes use of PIL library for picture loading, which is then transformed to a grayscale [0, 1] picture as a CuPy ndarray. For picture filtration we implement a sequence of img \u2192 R2C FFT \u2192 Gaussian filter \u2192 C2R iFFT \u2192 filtered_img. The Gaussian filter is (scriptstyle G(x,y)=expleft(-frac{x^2+y^2}{2sigma^2}proper)), which in frequency area can be a Gaussian (scriptstyle H(f_x,f_y)=expleft(-2pi^2sigma^2(f_x^2+f_y^2)proper)).<\/p>\n<\/div>\n<\/div>\n The next code exhibits the right way to apply a Gaussian picture filter with nvmath-python FFT and customized callback operate:<\/p>\n img = cp.asarray(Picture.open(“your_lovely_dog.jpg”).convert(“L”)) \/ 255.0 # Grey[0,1] # Implement Gaussian filter within the frequency area # Implement FFT epilog wrapper with the pre-defined signature # Compile epilog to LTO-IR concentrating on the present CUDA system # Compute R2C FFT utilizing nvmath-python with the compiled epilog # Compute C2R inverse FFT utilizing nvmath-python The second generally used state of affairs is looking nvmath-python system APIs from inside GPU kernels written in numba-cuda. nvmath-python helps system APIs for FFTs, GEMM, dense direct solvers (LU, Cholesky, QR) and RNG. The next instance exhibits the implementation of the Geometric Brownian Movement (GBM) for Monte Carlo inventory value simulations. It makes use of nvmath-python\u2019s random quantity generator for Gaussian distribution together with the customized numba-cuda code that converts regular distribution to the GBM Monte Carlo paths:<\/p>\n # Pre-compile the RNGs into IR to make use of alongside different system code # GBM parameters # Arrange CUDA kernel launch configuration # RNG initialization kernel # GBM path era kernel # Devour 4 regular variates at a time for higher throughput # Initialize RNG # Generate GBM paths on GPU Each operation in generate_gbm_paths has low arithmetic depth, which makes the host API-based implementation inefficient. It’s essential to get these operations fused with numba-cuda and nvmath-python system APIs.<\/p>\n Designed for productiveness with out efficiency compromises, nvmath-python reimagines the design of recent math libraries. Get began with one easy command:<\/p>\n![\"](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/image-84.webp\")
\nimport nvmath
\nimport cupy as cp<\/p>\n
\nwh = img.form[0] * picture.form[1] # We should normalize by the picture space
\nsigma_value = 20.0 # Filter dimension<\/p>\n
\ndef gaussian_filter(form, sigma):
\n fy = cp.fft.fftfreq(form[0])[:,None] # Column
\n fx = cp.fft.rfftfreq(form[1])[None,:] # Row
\n return = cp.exp(-2.0 * cp.pi * cp.pi * sigma * sigma * (fx * fx + fy * fy))<\/p>\n
\ndef epilog_impl(data_out, offset, knowledge, filter_data, unused): # Epilog to be compiled
\n data_out[offset] = knowledge * filter_data[offset] \/ wh<\/p>\n
\nepilog = nvmath.fft.compile_epilog(epilog_impl, “complex64”, “complex64”)<\/p>\n
\nh_filter = gaussian_filter(img.form, sigma)
\nimg_fft = nvmath.fft.rfft(picture, epilog={“ltoir”: epilog, “knowledge”: h_filter.knowledge.ptr})<\/p>\n
\nfiltered_img = nvmath.fft.irfft(img_fft) # Visualize or save as you need\n<\/p><\/div>\nCustomized numba-cuda kernels with nvmath-python calls<\/h3>\n
\nfrom nvmath.system import random
\nimport cupy as cp
\nimport math<\/p>\n
\ncompiled_rng = random.Compile(cc=None)<\/p>\n
\nrng_seed = 7777
\nn_time_steps, n_paths = 252, 8192
\nmu, sigma, s0 = 0.003, 0.027, 100.0<\/p>\n
\nthreads_per_block = 32
\nblocks = n_paths \/\/ threads_per_block + bool(n_paths % threads_per_block)
\nnthreads = threads_per_block * blocks<\/p>\n
\n@cuda.jit(hyperlink=compiled_rng.recordsdata, extensions=compiled_rng.extension)
\ndef init_rng(states, seed):
\n idx = cuda.grid(1)
\n random.init(seed, idx, 0, states[idx])<\/p>\n
\n@cuda.jit(hyperlink=compiled_rng.recordsdata, extensions=compiled_rng.extension)
\ndef generate_gbm_paths(states, paths, nsteps, mu, sigma, s0):
\n idx = cuda.grid(1)
\n if idx >= paths.form[0]:
\n return
\n paths[idx, 0] = s0<\/p>\n
\n for i in vary(1, nsteps, 4):
\n v = random.normal4(states[idx]) # Returned as float32x4 sort
\n vals = v.x, v.y, v.z, v.w # Decompose right into a tuple of float32
\n for j in vary(i, min(i + 4, nsteps)): # Course of a bit of 4 time steps
\n paths[idx, j] = paths[idx, j – 1] * math.exp(mu + sigma * vals[j – i])<\/p>\n
\nstates = random.StatesPhilox4_32_10(nthreads)
\ninit_rng[blocks, threads_per_block](states, rng_seed)<\/p>\n
\npaths = cp.empty((n_paths, n_time_steps), dtype=cp.float32, order=’F’)
\ngenerate_gbm_paths[blocks, threads_per_block](states, paths, n_time_steps, mu, sigma, s0)\n<\/p><\/div>\nGet began with nvmath-python<\/h2>\n