{"slug": "run-high-performance-core-math-at-scale-with-nvidia-nvmath-python", "title": "Run High-Performance Core Math at Scale with NVIDIA nvmath-python", "summary": "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.", "body_md": "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.\n\n## nvmath-python v1.0 release\n\nWith 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.\n\n## Fast and flexible installation\n\nInstalling 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.\n\n- Choose a package manager, such as pip, conda, uv, or pixi.\n- 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.\n- Pick and choose the CPU backend, device APIs support, or distributed APIs.\n- Choose a companion array library to work with, such as NumPy, CuPy, or PyTorch (or all of them). See the detailed\n[installation guide](https://docs.nvidia.com/cuda/nvmath-python/latest/installation.html)for available options.\n\n## A useful complement to existing array libraries\n\nLike 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.\n\nInstead, 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.\n\nIn the following example, nvmath-python consumes NumPy arrays and the result is also a NumPy array.\n\n``` python\nimport numpy as np\nimport nvmath\n\nm, n, k = 10, 40, 100\na = np.random.randn(m, k)  # a is a NumPy array\nb = np.random.randn(k, n)  # b is a NumPy array\nc = nvmath.linalg.advanced.matmul(a, b)  # c is also a NumPy array\n```\n\n## Choice of memory and execution spaces\n\nThe 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:\n\n- GPU libraries such as\n[cuBLAS](https://developer.nvidia.com/cublas)and[cuFFT](https://developer.nvidia.com/cufft). - CPU libraries such as\n[NVPL](https://developer.nvidia.com/nvpl)for NVIDIA Grace or any ARM v8 CPUs and Intel MKL for x86 hosts. - Distributed libraries such as\n[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).\n\nThis support simplifies code migration between CPU and GPU and enables hybrid and distributed workflows that combine CPU and GPU execution.\n\nThe following code illustrates how nvmath-python supports multiple memory and execution spaces.\n\n``` python\nimport cupy as cp\nimport numpy as np\nimport nvmath\n\nN = 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```\n\nThe `fft`\n\nexecution space for each call is inferred from its input tensor, either `a_gpu`\n\nor `a_cpu`\n\n, 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.\n\n## Generic and specialized APIs\n\nThe 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).\n\nGeneric 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.\n\nTo 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.\n\nThe 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`\n\nsubmodules to keep them distinct from generic APIs.\n\n## Logging with nvmath-python\n\nThe 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).\n\nThe following example illustrates the data flow between memory and execution spaces (using the advanced `matmul`\n\n).\n\n``` python\nimport numpy as np\nimport nvmath\nimport logging\n\nlogging.basicConfig(level=logging.INFO,\n    format=\"%(asctime)s %(levelname)-8s %(message)s\", force=True)\nlogging.disable(logging.NOTSET)\n\nm, n, k = 8000, 2000, 4000\na_cpu = np.random.randn(m, k).astype(np.float32)\nb_cpu = np.random.randn(k, n).astype(np.float32)\nd_cpu = nvmath.linalg.advanced.matmul(a_cpu, b_cpu)\n```\n\nThe produced output will look like:\n\n```\n2025-09-18 14:53:32,166 INFO     = SPECIFICATION PHASE =\n2025-09-18 14:53:32,167 INFO     The data type of operand A is 'float32', and that of operand B is 'float32'.\n2025-09-18 14:53:32,168 INFO     The input operands' memory space is cpu, and the execution space is on device 0.\n...\n```\n\nTake 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`\n\nto illustrate data flow between memory and execution spaces.\n\n``` python\nimport numpy as np\nimport nvmath\nimport logging\n\nlogging.basicConfig(level=logging.INFO,\n    format=\"%(asctime)s %(levelname)-8s %(message)s\", force=True)\nlogging.disable(logging.NOTSET)\n\nN = 10000\ne_cpu = (np.random.randn(N) + 1j * np.random.randn(N)).astype(np.complex64)\nr_cpu = nvmath.fft.fft(e_cpu)\n```\n\nThe logging output looks like:\n\n```\n2025-09-18 15:46:22,295 INFO     The FFT type is C2C.\n2025-09-18 15:46:22,295 INFO     The input data type is complex64, and the result data type is complex64.\n2025-09-18 15:46:22,296 INFO     The specified FFT axes are (0,).\n2025-09-18 15:46:22,297 INFO     The input tensor's memory space is cpu, and the execution space is cpu, with device cpu.\n2025-09-18 15:46:22,298 INFO     The specified stream for the FFT ctor is None.\n...\n```\n\nNote 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`\n\nkeyword argument to an API.\n\n## Why composite operations matter\n\nAn 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):\n\n\\(\\scriptstyle \\mathbf{D}=\\alpha\\mathbf{A}\\cdot\\mathbf{B}+\\beta\\mathbf{C}\\)\n\nThe following code illustrates GEMM on tall-and-skinny matrices with CuPy and nvmath-python.\n\n``` python\nimport cupy as cp\nimport nvmath\n\nm, n, k = 10_000_000, 40, 10\na = cp.random.randn(m, k, dtype=cp.float32)\nb = cp.random.randn(k, n, dtype=cp.float32)\nc = cp.random.randn(m, n, dtype=cp.float32)\nalpha, beta = 1.5, 0.5\n\nd1 = alpha * cp.matmul(a, b) + beta * c # Multiple kernels\nd2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta) # Single kernel\n```\n\nFigure 1 shows that a fused composite operation brings measurable benefits compared to NumPy-like APIs.\n\nnvmath-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.\n\n## Amortizing preparation costs by using stateful APIs\n\nAll 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.\n\n**Performance note**\n\nNVIDIA 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.\n\nIn 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.\n\nThe following example illustrates the use of class-form API for `matmul`\n\n(with `RELU_BIAS`\n\nepilog) on a batch of the `batch_size`\n\nsize of matrices `a`\n\nand `b`\n\n, and biases `bias`\n\n. The result of the prior matrix multiplication is an operand in the next matrix multiplication, and there are `feed_count`\n\noperations. Besides *planning* it also performs an *autotuning* phase.\n\n**The code below illustrates using **stateful APIs with planning, autotuning, and execution as distinct phases.\n\n``` python\nimport nvmath\nfrom nvmath.linalg.advanced import MatmulEpilog\nimport cupy as cp\n\nfeed_count = 10  # The operation feed count.\n\nbatch_size = 1024\nm, n, k = 1024, 1024, 1024\n\na = cp.random.rand(batch_size, m, k, dtype=cp.float32)\nb = cp.random.rand(batch_size, k, n, dtype=cp.float32)\nbias = cp.random.rand(batch_size, m, 1, dtype=cp.float32)\n\nwith nvmath.linalg.advanced.Matmul(a, b) as mm:\n    # 1. Planning phase\n    mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),\n            epilog_inputs={\"bias\": bias})\n\n    # 2. Autotuning phase\n    mm.autotune(iterations=5)\n\n    # 3. Execution phase.\n    for i in range(feed_count):\n        d = mm.execute()\n        # The result of the previous MM is the operand `a` of the next MM, so use\n        # reset_operands_unchecked() to reset the `a` operand.\n        mm.reset_operands_unchecked(a=d)\n```\n\nFigure 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.\n\nFigure 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.\n\n**Deployment note**\n\nAnother 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),\n\n[, and](https://github.com/NVIDIA/nvmath-python/blob/main/examples/linalg/advanced/matmul/example15_manual_tuning.py)\n\n`example15_manual_tuning.py`\n\n[in the nvmath-python GitHub repo.](https://github.com/NVIDIA/nvmath-python/blob/main/examples/linalg/advanced/matmul/example16_reuse_algorithms.py)\n\n`example16_reuse_algorithms.py`\n\n## Custom kernels fused with nvmath-python\n\nnvmath-python integrates with Python compilers such as `numba-cuda`\n\n, enabling high-performance custom Python code to be compiled just in time (JIT) and used alongside nvmath-python operations.\n\n### Custom FFT callbacks\n\n**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.\n\n**Gaussian filter example**\n\nAs 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`\n\n. For image filtration we implement a chain of `img`\n\n→ R2C FFT → Gaussian filter → C2R iFFT → `filtered_img`\n\n. 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)\\).\n\nThe following code shows how to apply a Gaussian image filter with nvmath-python FFT and custom callback function:\n\n``` python\nfrom PIL import Image\nimport nvmath\nimport cupy as cp\n\nimg = cp.asarray(Image.open(\"your_lovely_dog.jpg\").convert(\"L\")) / 255.0 # Gray[0,1]\nwh = img.shape[0] * image.shape[1] # We must normalize by the image area\nsigma_value = 20.0  # Filter size\n\n# Implement Gaussian filter in the frequency domain\ndef gaussian_filter(shape, sigma):\n    fy = cp.fft.fftfreq(shape[0])[:,None] # Column\n    fx = cp.fft.rfftfreq(shape[1])[None,:] # Row\n    return = cp.exp(-2.0 * cp.pi * cp.pi * sigma * sigma * (fx * fx + fy * fy))\n\n# Implement FFT epilog wrapper with the pre-defined signature\ndef epilog_impl(data_out, offset, data, filter_data, unused): # Epilog to be compiled\n    data_out[offset] = data * filter_data[offset] / wh\n\n# Compile epilog to LTO-IR targeting the current CUDA device\nepilog = nvmath.fft.compile_epilog(epilog_impl, \"complex64\", \"complex64\")\n\n# Compute R2C FFT using nvmath-python with the compiled epilog\nh_filter = gaussian_filter(img.shape, sigma)\nimg_fft = nvmath.fft.rfft(image, epilog={\"ltoir\": epilog, \"data\": h_filter.data.ptr})\n\n# Compute C2R inverse FFT using nvmath-python\nfiltered_img = nvmath.fft.irfft(img_fft) # Visualize or save as you want\n```\n\n### Custom numba-cuda kernels with nvmath-python calls\n\nThe 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:\n\n``` python\nfrom numba import cuda\nfrom nvmath.device import random\nimport cupy as cp\nimport math\n\n# Pre-compile the RNGs into IR to use alongside other device code\ncompiled_rng = random.Compile(cc=None)\n\n# GBM parameters\nrng_seed = 7777\nn_time_steps, n_paths = 252, 8192\nmu, sigma, s0 = 0.003, 0.027, 100.0\n\n# Set up CUDA kernel launch configuration\nthreads_per_block = 32\nblocks = n_paths // threads_per_block + bool(n_paths % threads_per_block)\nnthreads = threads_per_block * blocks\n\n# RNG initialization kernel\n@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)\ndef init_rng(states, seed):\n    idx = cuda.grid(1)\n    random.init(seed, idx, 0, states[idx])\n\n# GBM path generation kernel\n@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)\ndef generate_gbm_paths(states, paths, nsteps, mu, sigma, s0):\n    idx = cuda.grid(1)\n    if idx >= paths.shape[0]:\n        return\n    paths[idx, 0] = s0\n\n    # Consume 4 normal variates at a time for better throughput\n    for i in range(1, nsteps, 4):\n        v = random.normal4(states[idx])  # Returned as float32x4 type\n        vals = v.x, v.y, v.z, v.w  # Decompose into a tuple of float32\n        for j in range(i, min(i + 4, nsteps)):  # Process a chunk of 4 time steps\n            paths[idx, j] = paths[idx, j - 1] * math.exp(mu + sigma * vals[j - i])\n\n# Initialize RNG\nstates = random.StatesPhilox4_32_10(nthreads)\ninit_rng[blocks, threads_per_block](states, rng_seed)\n\n# Generate GBM paths on GPU\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```\n\nEvery operation in `generate_gbm_paths`\n\nhas 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.\n\n## Get started with nvmath-python\n\nDesigned for productivity without performance compromises, nvmath-python reimagines the design of modern math libraries. Get started with one simple command:\n\n```\npip install nvmath-python[cu13]\n```\n\nAdditional resources include:\n\n[Installation documentation](https://docs.nvidia.com/cuda/nvmath-python/latest/installation.html)with detailed setup instructions- The\n[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,\n[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\n[nvmath-python product page](https://developer.nvidia.com/nvmath-python)for the latest news and announcements\n\n## Acknowledgments\n\nThe library is a result of efforts of many people from across NVIDIA, including:\n\n*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.*", "url": "https://wpnews.pro/news/run-high-performance-core-math-at-scale-with-nvidia-nvmath-python", "canonical_source": "https://developer.nvidia.com/blog/run-high-performance-core-math-at-scale-with-nvidia-nvmath-python/", "published_at": "2026-07-30 22:43:04+00:00", "updated_at": "2026-07-30 23:01:12.450397+00:00", "lang": "en", "topics": ["ai-infrastructure", "developer-tools"], "entities": ["NVIDIA", "nvmath-python", "CUDA-X", "cuFFT", "cuBLASLt", "cuDSS", "cuSPARSE", "cuTENSOR"], "alternates": {"html": "https://wpnews.pro/news/run-high-performance-core-math-at-scale-with-nvidia-nvmath-python", "markdown": "https://wpnews.pro/news/run-high-performance-core-math-at-scale-with-nvidia-nvmath-python.md", "text": "https://wpnews.pro/news/run-high-performance-core-math-at-scale-with-nvidia-nvmath-python.txt", "jsonld": "https://wpnews.pro/news/run-high-performance-core-math-at-scale-with-nvidia-nvmath-python.jsonld"}}