nvmath-python: plan, execute and fuse with NVIDIA math libraries

The knobs the array frameworks leave out

The stated aim is to give Python users access to the full feature set of NVIDIA's math libraries, and the differentiator is precision rather than convenience. The documentation says the API allows access to all parameters of the underlying NVIDIA cuBLASLt library, which it names as a distinguishing feature against other wrappings of the NVIDIA C API libraries. The package is designed to work with existing Python array and tensor frameworks and to focus on functionality that is missing from them, and it accepts input tensors from PyTorch, CuPy and NumPy.

The APIs fall into three groups, and the distinction is about where the code runs rather than what it computes. Host APIs are called from host code but can execute in any supported execution space, CPU or GPU. Distributed host APIs are called from host code but execute on a distributed multi-node, multi-GPU system. Device side APIs expose the on-device interfaces, so the code that would normally run on the host is compiled and launched from the device itself. The last group is the one to read about carefully, because it changes who owns the thread and the memory.

Planning once, executing many times

The host API example is built around a stateful object rather than a function call, and that is the design point. A `Matmul` object is constructed with the operands and an options dictionary, and the options are the fine tuning that a plain call cannot express, such as selecting a mixed-precision compute type:

mm = nvmath.linalg.advanced.Matmul(
    a,
    b,
    options={
        "compute_type": nvmath.linalg.advanced.MatmulComputeType.COMPUTE_32F_FAST_16F
    },
)

Planning is a separate step, and it returns a sequence of algorithms that can be configured. The plan call is also where an epilogue is attached, which is an operation applied to the result of the multiplication without a separate function call:

mm.plan(
    epilog=nvmath.linalg.advanced.MatmulEpilog.BIAS,
    epilog_inputs={"bias": bias},
)

Execution is then a separate call, which is what makes the object reusable:

# Execute the matrix multiplication.
result = mm.execute()

The object holds resources, and the documentation is explicit that it must be freed when finished or used as a context manager. The example also reminds you to synchronise the default stream, because execution is non-blocking for GPU operands by default, so a program that reads the result without synchronising can read it before it is ready.

FFT epilogues written in Python

The FFT side goes further than selecting a built-in option. nvmath-python lets you write custom prologues and epilogues for FFT functions as ordinary Python functions and compile them to LTO-IR, which is then handed to the library. The example defines a rescaling function that divides the output by the square root of the transform size, which is the normalisation a unitary transform needs:

def rescale(data_out, offset, data, user_info, unused):
    data_out[offset] = data * norm_factor

The function is compiled on a CUDA device context, with the input and output types named explicitly:

with cp.cuda.Device():
    epilog = nvmath.fft.compile_epilog(rescale, "complex128", "complex128")

The compiled epilogue is then passed into the transform, and the point of doing this at all is that the rescaling happens inside the same operation as the transform rather than in a second pass over the data. The example closes by comparing the fused result against a separate normalised transform with an assertion, which is the right way to check that a custom epilogue has not changed the answer.

Distributed FFT on more than one node

The distributed path is a different world, and the documentation is careful about the constraints. cuFFTMp uses the NVSHMEM partitioned global address space model, which requires GPU operands to sit on the symmetric heap, so the input buffer is allocated with `allocate_symmetric_memory` rather than as an ordinary CuPy array, and it is then filled inside a device context. Initialisation takes a device id, the MPI communicator and a list of backends:

nvmath.distributed.initialize(device_id, comm, backends=["nvshmem"])

The distribution itself is explicit. A three dimensional transform is spread over the processes with the `Slab` distribution on one axis, the forward transform can be told not to redistribute the result, and the inverse transform then runs with the opposite slab assignment, so the buffers change shape between steps. The transforms are in place, which means the input, output and inverse output share one symmetric memory buffer, and the operand that is not garbage collected has to be freed by hand. That free is a collective operation that all processes must call at the same point in the execution, and the example comments make the point explicitly.

What the build actually requires

This is not a pure Python package. The build system requires Cython 3.2.0 or newer, setuptools 77.0.3 or newer, tomli on Python older than 3.11, and the CUDA toolkit with the cudart and nvcc extras, bounded between 13.1 and 14. The extension modules are C++ extensions generated from `.pyx` sources, and `setup.py` builds the module list by reading the module names out of the project configuration, with a helper that can rename a module, insert a submodule into its dotted name and pick a platform specific source suffix such as `_linux` or `_windows`. At runtime the package depends on `cuda-bindings`, `cuda-core` between 0.5.0 and 2, `cuda-pathfinder` between 1.5.1 and 2.0, NumPy between 1.25 and 3, and `pywin32` on Windows. The supported Python range is 3.10 up to but not including 3.15, and the classifier calls the project production stable at version 1.0.0.

Where it is the wrong tool

Three situations call for something else. If your work is ordinary array manipulation, the frameworks you already use cover the common cases and this package adds a build dependency to get at features you will not use. If your environment cannot supply the CUDA toolkit and nvcc, the package cannot be built at all, which rules it out for a CPU-only pipeline or a machine that must avoid vendor specific tooling. And if you only need a distributed transform, the same job may be simpler through the framework's own distributed layer, without taking on symmetric heap semantics and collective frees.

On the project itself, the release history shows a fast climb to stability, with v0.8.0 in January 2026, v0.9.0 in April and v1.0.0 in July, and the last push was on 2026-07-14, the same day as the 1.0.0 release. The licence is Apache-2.0. The repository keeps the package source in `nvmath/`, the extension builder in `builder/`, per-area examples in `examples/` covering device, distributed, fft, linalg, sparse and tensor code, notebooks for documentation, release notes managed with reno, and a tests directory. The documentation lives with NVIDIA's CUDA documentation, not in the repository.

Editorial conclusion

nvmath-python fits a performance engineer who has already outgrown the high level calls in CuPy or PyTorch and needs the library knobs exposed, since the plan and execute split, the fused epilogues and the LTO-IR path are exactly the parts a thin wrapper normally hides. It is the wrong choice for ordinary array work, since the frameworks already cover the common cases, and it is the wrong choice if your build environment cannot provide the CUDA toolkit, because the build requirements include nvcc and the CUDA runtime. Before adopting it, check that your Python version falls inside the supported range of 3.10 to 3.14, pin a version rather than tracking the branch, since the 1.0.0 release is recent and the 0.8 to 0.9 line moved quickly, and read the distributed examples carefully if you plan to use the multi-node path, because the symmetric heap buffers are not garbage collected and freeing them is a collective operation that every process must perform at the same point.

Frequently asked questions

Official sources

1. License: Apache-2.0
2. NVIDIA/nvmath-python on GitHub
3. Project website
4. README
5. Releases