NVIDIA: Rustba ültették át a CUDA Tile kerneljeit egy AI-ügynökkel

Közös köztes reprezentációra épül az átültetés

Az NVIDIA szeptember 16-án bemutatott megoldása a cuTile Rustot használja a Rust nyelvű, csempékre épülő GPU-kernel-fejlesztéshez. A rendszer kiterjeszti a Rust tulajdonlási modelljét a tile-alapú GPU-kernelekre: a módosítható kimeneteket egymástól független részekre bontja, miközben megőrzi a gazdagépen érvényes tulajdonlási szerződést a kernelindítások között.

A CUDA Tile család három előtere, a cuTile Python, a Triton-TileIR és a cuTile Rust ugyanazt a CUDA Tile IR köztes reprezentációt használja. Ezekből a tileiras fordító állítja elő a GPU binárisát. Az NVIDIA szerint ez azért teszi praktikussá az átültetést, mert a Rustra portolás ugyanannak a tile-programnak a biztonságosabb gazdanyelven való újrafogalmazása, nem pedig újraoptimalizálási feladat.

A közös IR az ellenőrzést is lehetővé teszi. A fejlesztők összehasonlíthatják a kiinduló és az átültetett kernel IR-jét, többek között a memória-műveletek családját, a tile-ok alakját és a redukciókat. Így már a funkcionális tesztelés előtt kiszűrhetők például a hibás költségjavaslatok vagy a hiányzó oszthatósági attribútumok.

Az AI-ügynök minden lépésnél ellenőriz

A TileGymben található AI-ügynöki készség a cuTile Python- és Triton-TileIR-kerneleket cuTile Rust-kódra fordítja. Az NVIDIA csapata ezzel mind a 24 nyilvános TileGym-operátort átültette. A nagyjából 40 GPU-kernel között elemi műveletek, flash-attention decode, Multi-head Latent Attention és mixture-of-experts modellekhez tartozó megoldások is vannak.

A fordítás korlátozott, többügynökös munkafolyamatban zajlik. A folyamat elemzést, eszközoldali kernelt, gazdaoldali és FFI-kódot, valamint teljesítménymérést tartalmaz. Minden szakasz végén géppel ellenőrizhető döntés születik, a továbblépést validátorscriptek és a Tile IR összehasonlítása határozza meg.

A rendszer egyik központi feladata a két nyelv specializációs különbségének kezelése. A cuTile Python a hívás pillanatában, implicit módon specializálja a kernelt. Rustban minden specializációt deklarálni kell a kernel aláírásában, ezért a típus- és alakmegkötéseket a rustc fordító ellenőrizheti. Egy Python-kernelből emiatt több strukturális Rust-belépési pont is készülhet, ha egy feltétel megváltoztatja például a tile rangját.

A Rust-forrás előre fordul, a konkrét konstansgenerikus értékekkel történő futásidejű specializáció és a cubin létrehozása pedig az első indításkor történik. A GPU-binárist a rendszer ezt követően gyorsítótárazza.

A teljesítmény átlagosan 99,5 százalékos

Az NVIDIA DGX B200 rendszeren mérte meg az átültetett operátorokat. A közlemény szerint mind a 24 operátor teljesítette a cuTile Pythonhoz viszonyított 0,95-ös geometriai átlagú gyorsulási küszöböt, a teljes eredmény pedig 0,995, vagyis a cuTile Python teljesítményének 99,5 százaléka lett.

A cuTile Rust a TileGymben egy további háttérként használható. A tilegym.set_backend("cutile-rs") beállítás ugyanazt az operátor-API-t irányítja a Rust-kernelekhez. A C-ABI-réteg tensorleírókat ad át másolás vagy új memóriafoglalás nélkül, így a fordítás késleltetetten történhet, és a fejlesztési ciklus is közvetlenül használhatja az átültetett kerneleket.

A megoldás a fejlesztők számára azért lehet fontos, mert a TileGym meglévő, Pythonban és Triton-TileIR-ben írt kernelkönyvtára Rustból is elérhetővé válik. Az NVIDIA az ügynöki készséget a TileGym-tárolóban teszi elérhetővé, ahol a saját kernelek fordításához szükséges beállításokat és használati útmutatót is meg lehet találni.