nvmath-python Device APIs#

The device module of nvmath-pythonnvmath.device offers integration with NVIDIA’shigh-performance computing libraries through device APIs for cuFFTDx, cuBLASDx, and cuRAND.Detailed documentation for these libraries can be found atcuFFTDx,cuBLASDx, andcuRAND device APIs respectively.Device APIs can only be called from CUDA device or kernel code, and execute on the GPU.

Users may take advantage of the device module via the two approaches below:

  • Numba Extensions: Users can access these device APIs via Numba by utilizing specificextensions that simplify the process of defining functions, querying device traits, andcalling device functions.

  • Third-party JIT Compilers: The APIs are also available through low-level interfaces inother JIT compilers, allowing advanced users to work directly with the raw device code.

Note

TheFFT andMatmul device APIs in modulenvmath.device currently supports cuFFTDx 1.5.1 and cuBLASDx 0.4.1, also availableas part of MathDx 25.06. All functionalities from the C++ libraries are supported withthe exception of cuFFTDx C++ APIs with a workspace argument, which are currently notavailable in nvmath-python.