Az nvmath-python v1.0 általános elérhetőségével az NVIDIA egy Python‑barát réteget tett elérhetővé a CUDA‑X matematika‑könyvtárak (például cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR, cuBLASMp és társai) fölé. A cél az, hogy a Python tudományos és mérnöki közössége hozzáférhessen a CUDA‑X teljesítményéhez anélkül, hogy meglévő Python‑alapú munkafolyamatokat kellene átírni C/C++‑ra. Az nvmath-python képes CPU‑n, egyetlen CUDA‑képességű GPU‑n, vagy kiterjedt, több‑GPU, több‑node elrendezéseken futtatni műveleteket a választott API‑tól függően.
Mit kínál a könyvtár
nvmath-python egy Python‑oszlopos absztrakció, amely két fő célt szolgál: (1) általános numerikus műveleteket biztosít, amelyekre számos tudományos és mérnöki alkalmazásnak szüksége van, és (2) közvetlen hozzáférést ad a CUDA‑X‑ben meglévő, nagy optimalizáltságú rutinok teljes funkcionalitásához. Nem célja helyettesíteni a teljes értékű tömbkönyvtárakat (például NumPy‑t vagy CuPy‑t) hagyományos tömbműveletek‑indexelés/szeletelés/redukciók terén; ehelyett az optimalizált GPU‑ és elosztott rutinok használatát teszi egyszerűbbé Pythonból.
A könyvtár különlegességei közé tartozik a "universal sparse tensor" (UST) megközelítés, amely domain‑specifikus nyelven engedi a felhasználónak saját, alkalmazás‑optimalizált ritkított formátumot definiálni anélkül, hogy azt alacsony szinten kellene implementálni.
Telepítés és konfigurálhatóság
Az nvmath-python célja, hogy a natív függőségeket tartalmazó Python csomagok telepítését gyorsabbá és rugalmasabbá tegye. A felhasználó választhat pip, conda, uv vagy pixi csomagkezelőt. Lehetőség van teljes függőség‑feloldással telepíteni minden szükséges natív komponenst, vagy minimális (bare‑minimum) telepítést végezni, ami hasznos például CI/CD vagy CPU‑csak környezetekben.
A telepítéskor kiválasztható a CPU‑backend (például NVPL vagy Intel MKL), a kívánt eszköz‑API‑k és az elosztott API‑k támogatása, valamint a társtömbkönyvtár: NumPy, CuPy, PyTorch vagy ezek tetszőleges kombinációja. A gyors indításhoz a példa telepítőparancs:
pip install nvmath-python[cu13]
Tömbkönyvtárakkal való együttműködés
nvmath-python úgy működik, hogy meglévő tömbkönyvtárakat fogyaszt és eredményül is azokat adja vissza — például egy NumPy‑t bemenetként használó hívás eredménye továbbra is NumPy tömb lesz. Ez megkönnyíti a kód migrálását CPU és GPU között, illetve a hibrid/elosztott munkafolyamatokat.
Példa: NumPy‑bemennelű mátrixszorzás eredménye NumPy‑ként tér vissza.
import numpy as np
import nvmath
m, n, k = 10, 40, 100
a = np.random.randn(m, k)
b = np.random.randn(k, n)
c = nvmath.linalg.advanced.matmul(a, b)
Memória‑ és végrehajtási terek kiválasztása
A könyvtár képes kezelni különböző memória‑ és végrehajtási tereket: CPU oldali (pl. NumPy + NVPL/MKL), GPU oldali (pl. CuPy + cuBLAS/cuFFT), illetve elosztott könyvtárak (pl. cuBLASMp, cuSOLVERMp, cuFFTMp). Az nvmath‑python az input tömb alapján is képes automatikusan kiválasztani az FFT vagy más művelet végrehajtási terét, de a felhasználó manuálisan is megadhat egy execution paramétert.
Példa CPU és GPU FFT‑hívásra:
import cupy as cp
import numpy as np
import nvmath
N = 2048
a_gpu = cp.random.randn(N) + 1j * cp.random.randn(N)
a_cpu = np.random.randn(N) + 1j * np.random.randn(N)
c_gpu = nvmath.fft.fft(a_gpu)
c_cpu = nvmath.fft.fft(a_cpu)
A könyvtár beépített naplózási képességeivel látható, hogy melyik művelet hol futott, és honnan kellett adatot átmásolni.
Generikus és specializált API‑k
Az nvmath-python API‑i két osztályba sorolhatók: generikus API‑k (széles, de felszínes funkcionalitás) és specializált API‑k (szűk, de mély konfigurálhatóság). A generikus API‑k egységes felhasználói élményt kínálnak több memória‑ és végrehajtási térre, míg a specializált API‑k (az nvmath.linalg.advanced almodulban) teljes hozzáférést adnak hardver‑közeli optimalizációkhoz, például egy GPU‑specifikus, sűrű mátrixokra tervezett, összetett mátrixszorzás minden konfigurációs lehetőségéhez.
A felhasználási eset dönti el, hogy mikor érdemes a specializált (teljesítmény‑kritikus) vagy a generikus (kényelmesebb) API‑kat választani.
Naplózás és adatmigrációk láthatósága
A könyvtár integrálódik a Python standard logging modullal és részletes információkat ad a specifikációs és végrehajtási fázisokban. Például egy 8000×4000‑as matmul futtatásakor a logok megmutatják az operandusok adattípusát, memória‑térét és azt is, hogy az adott művelet mely eszközön (device 0, CPU stb.) futott. A napló információi segítenek azonosítani a költséges adatátviteli műveleteket.
A példa logbejegyzések (az eredetiben 2025‑09‑18 dátummal szerepelnek) tanúsága szerint az nvmath‑python, amikor csak lehet, úgy választ végrehajtási teret, hogy minimalizálja az adatmozgásokat; azonban a felhasználó mindig felülírhatja ezt.
Összetett műveletek és kernel‑fúzió
Az olyan összetett műveletek, mint D = f(α A B + β C), tipikusan alacsony aritmetikai intenzitású lépéseket tartalmaznak, amikor a primitív műveletek egymás után futnak. Az nvmath‑python a mögöttes cuBLASLt‑nek köszönhetően képes JIT kernel‑fúzióra, azaz egyetlen kernelben végrehajtani az összetett műveletet, ami különösen előnyös például "tall‑and‑skinny" mátrixokkal végzett GEMM esetén.
Példa: egy nagyon hosszú (m = 10_000_000), keskeny (k = 10) mátrixszorzásnál a NumPy/CuPy stílusú láncolt meghívás több kernel‑indítást és adatmozgást eredményez, míg az nvmath.linalg.advanced.matmul egyetlen kernelben futtatva jobb teljesítményt ad.
Állapotmentes (stateless) vs. állapotfüggő (stateful) API‑k: tervezés és autotuning
A könyvtár két használati modellt támogat: egyszeri hívásos, állapotmentes API‑kat és osztályalapú, állapotfüggő API‑kat. Az állapotmentes hívások kényelmesek, de minden egyes végrehajtáskor fut egy előkészítési (planning) fázis, amely időigényes lehet, különösen ha autotuning is történik. Az állapotfüggő API lehetővé teszi a tervek létrehozását, autotuningolását és többszöri végrehajtását, így a plan/autotune költsége több futáson eloszlik.
Példa: egy Matmul osztály használatával a felhasználó külön tervezi a műveletet (plan), futtat egy autotuning fázist (autotune), majd többször végrehajtja (execute). Az autotuning iterációk száma konfigurálható; az elkészült, autotuned terv pedig lementhető és más munkamenetekben újrahasználható.
A teljesítményjegyzet szerint az autotuning sokszor hasznos lehet, de nem mindig szükséges: bizonyos GPU‑kon (például az egyik méréssorozatban az NVIDIA RTX A6000) az autotuning nagy, akár 256%-os gyorsulást eredményezett egy adott probléma esetén, míg más hardverek (pl. NVIDIA B200 a bemutatott példában) már alapértelmezett heuristics mellett is csúcsteljesítményt adtak.
Egyedi JIT‑komponensek és numba‑cuda integráció
nvmath-python támogatja a JIT‑kompilált, felhasználó által definiált callbackeket és eszköz‑API hívásokat. A numba‑cuda‑val együtt használva lehetséges saját GPU‑kernelt írni, amely közben nvmath‑python device API‑kat (például FFT, GEMM, direkt sűrű megoldók, RNG) hív meg.
Példák a dokumentációból:
- Egy JIT‑kompilált FFT‑epilog (vagy prolog) Python függvényként írható meg, majd LTO‑IR‑re fordítható, és az FFT végrehajtás közben alkalmazható egy egyéni szűrő (például Gauss‑filter képfeldolgozásra).
- numba‑cuda kernelből meghívott nvmath‑python eszközök példája a Geometric Brownian Motion (GBM) Monte Carlo szimuláció, ahol a GPU‑oldali RNG és a GPU‑kernel együttműködve generálja a pályákat.
Ezek a lehetőségek különösen fontosak, amikor az egyes operációk alacsony aritmetikai intenzitásúak, és kritikus, hogy a hívások össze legyenek fuzionálva a GPU‑kódon belül.
Teljesítmény‑ és üzembehelyezési megfontolások
- Autotuning: az nvmath‑python és a mögöttes CUDA‑X könyvtárak heuristikát és autotuningot használnak a legjobb kernel kiválasztásához. Az autotuning időigényes lehet, de gyakran jelentős gyorsulást hoz. Az autotuned tervek sorosíthatók és több, homogén rendszeren is újrahasználhatók.
- Fúzió és kompozit műveletek: kernel‑fúzió növeli az aritmetikai intenzitást, és így jobb hatékonyságot biztosít alacsony intenzitású összetett műveleteknél.
A cikk bemutatását kiegészítő példák az nvmath‑python GitHub repójában találhatók (például example14_autotune.py, example15_manual_tuning.py, example16_reuse_algorithms.py).
Hogyan kezdjünk hozzá
A gyors kipróbáláshoz elegendő egy parancs:
pip install nvmath-python[cu13]
További anyagok: telepítési dokumentáció, GitHub repó példákkal és jegyzetfüzetekkel, NVIDIA Accelerated Computing Hub‑on bővített Python oktatási anyagok, valamint egy korábbi bejegyzés a universal sparse tensorról.
Elismerések
A projekt számos NVIDIA‑s munkatárs hozzájárulásának eredménye. A név szerint említettek között szerepelnek többek között 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 és mások.
