Run High-Performance Core Math at Scale with NVIDIA nvmath-python

Summary

NVIDIA released nvmath-python v1.0 to give Python users direct access to CUDA-X math libraries without requiring low-level C or C++ interfaces. The library acts as an abstraction layer over GPU engines like cuFFT and cuBLASLt, CPU backends like NVPL and Intel MKL, and distributed libraries such as cuBLASMp. It interfaces directly with NumPy, CuPy, and PyTorch tensors, structuring routines into generic interfaces for broad interoperability and specialized advanced modules for hardware-specific optimizations. Users can define custom sparse formats with Universal Sparse Tensors or invoke compiled device APIs directly inside custom numba-cuda kernels. Built-in logging captures execution space decisions and flags potentially expensive memory transfers across host and device environments.

Takeaways

1. 1

   nvmath-python exposes CUDA-X and NVPL libraries across CPUs, GPUs, and distributed clusters while directly accepting and returning NumPy, CuPy, and PyTorch arrays.

2. 2

   APIs are separated into generic multitools that prioritize broad compatibility across execution spaces and specialized advanced submodules that deliver hardware-tuned composite operations.

3. 3

   Low-arithmetic-intensity operations can be fused by invoking compiled nvmath-python device APIs, such as Philox RNG or direct solvers, directly inside custom numba-cuda kernels.