NVIDIA nvmath-python v1.0 Brings CUDA-X Math to Python at Scale

NVIDIA released nvmath-python v1.0, a Pythonic abstraction layer over CUDA-X math libraries including cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR, and cuBLASMp. The library targets Python users who need CUDA-X performance for common math operations without leaving their existing workflows. Depending on the API, operations can run on a CPU, a CUDA-enabled GPU, or distributed multi-GPU, multi-node systems.

Installation and array library flexibility

Installation supports pip, conda, uv, or pixi, with options for full dependency resolution or bare-minimum installs for scenarios like CI/CD or CPU-only environments. Users can select CPU backends, device API support, distributed APIs, and a companion array library such as NumPy, CuPy, or PyTorch. The library is not meant to replace general-purpose array libraries, so it does not provide indexing, slicing, or reduction. Instead, it exposes the full functionality of CUDA-X math libraries in Python, letting existing frameworks use GPU-accelerated routines without low-level C/C++ interfaces.

The library consumes and returns arrays from the chosen companion library. For example, a matrix multiplication can take NumPy arrays and return a NumPy array. The execution space for each call is inferred from the input tensor, so an FFT on a CuPy array runs on the GPU while the same call on a NumPy array runs on the CPU. Users can also specify a different execution space explicitly. The library logs where each operation ran, which helps identify potentially expensive data transfers between memory and execution spaces.

Generic and specialized APIs

nvmath-python divides its APIs into two classes. Generic APIs provide a uniform experience across execution spaces, memory spaces, and operand types, but restrict configurability to baseline features shared across their broad scope. Specialized APIs offer a comprehensive set of features for a narrow operational range and may be restricted to particular hardware. All specialized APIs live in advanced submodules to keep them distinct.

The advanced matrix multiplication implements the composite operation D equals f(A times B plus C) specifically for dense operands on the GPU, with every configuration needed to maximize hardware efficiency. The generic matrix multiplication API accommodates dense and structured operands across CPU and GPU execution spaces, but offers only the common subset of options. Specialized APIs are ideal when an operation becomes a computational bottleneck, while generic APIs suit tasks that are not performance-critical.

Composite operations and stateful APIs

Chaining primitive operations as a series of calls is inefficient when those operations have low arithmetic intensity. A notable example is GEMM with a tall-and-skinny matrix, where the fused composite operation in nvmath-python runs in a single kernel thanks to cuBLASLt's just-in-time kernel fusion. This brings measurable benefits compared to NumPy-like APIs that require multiple kernels.

The functional-form, or stateless, API involves a time-consuming preparation phase that includes planning and autotuning. Stateful APIs amortize these costs across repeated executions. Autotuned plans can be serialized to disk for reuse across sessions or deployment on homogeneous systems. The library also integrates with numba-cuda for JIT-compiled custom kernels and callbacks for FFT, GEMM, dense direct solvers, and random number generation, enabling fusion of low-arithmetic-intensity operations.

Logging integrates with the Python standard library logger, capturing computational details at debug, information, warning, and error levels. The logs show the data flow between memory and execution spaces, including the data type of operands, the execution device, and the specified stream. Whenever possible, nvmath-python selects the execution space to minimize data transfer overheads, but users can override this by providing an execution keyword argument.