Tým NVIDIA podle svého technického blogu vytvořil agentní dovednost (skill), která překládá GPU kernely z cuTile Python a Triton-TileIR do cuTile Rust. Autoři s její pomocí převedli všech 24 veřejných operátorů knihovny TileGym. V průměru dosáhli 99,5 % výkonu verze v cuTile Python. Výkon měřili na stroji NVIDIA DGX B200 na 347 spárovaných konfiguracích. Postup se opírá o společný mezijazyk CUDA Tile IR, díky němuž lze výstup agenta porovnat s referenčním kernelem strukturálně, nikoli jen testy. Následující text shrnuje, jak překlad funguje, v čem se Python a Rust liší, jak je řízena práce subagentů a jaká omezení autoři sami uvádějí.
Co je cuTile Rust
cuTile Rust, zkráceně cutile-rs, je podle NVIDIA systém pro bezpečné a idiomatické psaní GPU kernelů v jazyce Rust. Rozšiřuje vlastnický model Rustu na dlaždicové kernely. Měnitelné výstupy rozděluje na disjunktní části, takže se jednotlivé části výstupu nepřekrývají. Napříč spuštěními kernelů navíc zachovává vlastnický kontrakt na straně hostitele. Právě do tohoto systému agentní dovednost přenáší kernely, které dosud existovaly v cuTile Python a Triton-TileIR. Autoři tím chtějí zpřístupnit všechny tyto kernely i v Rustu.
Rozsah převedené sady
Převáděné operátory obsahují dohromady zhruba 40 GPU kernelů. Podle autorů sahají od element-wise operací přes flash-attention decode a Multi-head Latent Attention (MLA) až po modely mixture-of-experts (MoE). Některé operátory přitom vyžadují více variant kernelu, takže počet kernelů je vyšší než počet operátorů. Autoři uvádějí, že převedli všech 24 veřejných operátorů knihovny TileGym.
Společný základ: jeden mezijazyk pro tři front-endy
Celý postup stojí na jednom pozorování. cuTile Python, Triton-TileIR a cuTile Rust jsou podle autorů tři front-endy nad stejným mezijazykem CUDA Tile IR, tedy dialektem cuda_tile. Všechny tři předávají výstup témuž kompilátoru tileiras. Ten provádí optimalizace na úrovni dlaždic a generuje GPU binárku. Z toho autoři vyvozují, že portování kernelu do cuTile Rust není problém re-optimalizace. Jde o přepsání stejného dlaždicového programu do bezpečnějšího hostitelského jazyka. Pod ním zůstává stejný kompilátor a stejný výkonový model. Překlad se tedy podle autorů nemusí začínat novými optimalizacemi, protože všechny tři front-endy se sbíhají v témže mezijazyku.
Strukturální kontrola místo samotných testů
Sdílený mezijazyk umožňuje překlad kontrolovat strukturálně. Protože všechny tři front-endy vydávají stejný dialekt, lze před spuštěním jediného testu vypsat Tile IR referenčního kernelu i Tile IR přeloženého kernelu a oba výpisy porovnat (diff). Výstup agenta se tak ověřuje nejen funkčně, ale i podle struktury. Autoři upozorňují, že věrohodně vypadající, ale chybný překlad může testy projít a přesto být nesprávný mimo oblast, kterou testy pokrývají. Může také přinést zhoršení výkonu. Strukturální porovnání takové chyby zachytí. Uvádějí například TMA načtení se špatným cost hintem nebo ztracený atribut dělitelnosti.
Kde se Python a Rust rozcházejí
Hlavní obtíž překladu spočívá podle autorů jinde než ve výkonu. JIT kompilace v cuTile Python specializuje každý kernel implicitně v okamžiku volání. Rust naproti tomu vyžaduje, aby byla každá specializace deklarovaná v signatuře kernelu. Zdroj v Rustu se kompiluje předem a rustc kontroluje tvary dlaždic i typy prvků. Crate přitom obsahuje syntaktický strom kernelu. Při prvním spuštění ho runtime specializuje konkrétními hodnotami konstantních generik a zkompiluje cubin, který se poté ukládá do cache. GPU binárka se tedy stále kompiluje JIT, ale implicitnost zmizí: nic se nespecializuje, pokud to signatura kernelu výslovně neuvádí.
Rozdíl má konkrétní důsledky pro překlad. Nevyužité větve s ct.Constant se v Pythonu před kompilací zahodí. V Rustu musí typovou kontrolou projít obě větve. Jeden Python kernel se proto může rozpadnout na několik strukturních vstupů v Rustu. Autoři jako příklad uvádějí kernel layer_norm, který se dělí na 2D a 1D variantu, protože větev mění rank dlaždice.
Napojení na TileGym přes vrstvu C-ABI
Kernely se do TileGym napojují přes vrstvu C-ABI. Tenzory se přes ni předávají jako jednoduchý deskriptor s ukazatelem, počtem dimenzí, tvarem a kroky. Podle autorů se přitom nic nekopíruje, nealokuje ani nepřebírá vlastnictví. Rust proto nikdy neuvolňuje paměť, kterou nevlastní. Kernel se spouští asynchronně na CUDA streamu volajícího.
Praktická výhoda se týká vývoje. Protože se cuTile Rust kompiluje líně, backend sleduje aktuálnost zdrojů. Po úpravě souboru kernel.rs nebo manifestu crate se sdílená knihovna při dalším volání sama přebuduje. Explicitní cargo build ve vývojovém cyklu není potřeba. Autoři proto tvrdí, že iterace nad kernelem v Rustu je tak snadná jako v Pythonu.
Orchestrátor, který sám nepracuje
Dovednost se jmenuje tilegym-converting-python-to-rust a funguje jako orchestrátor. Agent, který ji načte, sám nedělá žádnou inženýrskou práci a pouze směruje úkoly. Jeho jedinou pravomocí je směrování. Práci provádějí specializovaní subagenti: analyzátor, autor kernelu, builder hostitelského kódu a FFI a validátor výkonu. K nim se přidávají dva diagnostičtí specialisté. Zapojí se jen tehdy, když něco selže.
Pevná pravidla řízení
Řízení běhu se podle autorů opírá o pevná pravidla. Subagenti komunikují jen přes artefakty s pevným schématem, nikoli prostřednictvím konverzace. Každá fáze končí strojově kontrolovatelným verdiktem. Orchestrátor se rozhoduje podle verdiktu, nikoli podle prózy, a opravu nikdy nevyvozuje z volného textu. Počty pokusů jsou pevně omezené, například na dva pokusy autora kernelu a dva pokusy hostitelského builderu. Pipeline tak podle autorů buď skončí v rámci rozpočtu, nebo se zastaví s diagnózou. Jednotlivé kroky ukazuje vývojový diagram.
Co ukázalo měření
Výkon autoři měřili pomocí CUPTI na stroji NVIDIA DGX B200, kde měl každý backend vyhrazený jeden GPU. Porovnávali 347 spárovaných konfigurací napříč 24 operátory. U každé konfigurace použili nejlepší měření ze čtyř běhů CI. Celkový geometrický průměr dosáhl podle autorů hodnoty 0,995, což označují za paritu s cuTile Python. Odpovídá to údaji o 99,5 % výkonu verze v Pythonu.
Další údaje autoři uvádějí u jednotlivých operátorů. Všech 24 operátorů splnilo práh 0,95 oproti cuTile Python. Asi třetina operátorů referenci překonala. Největší zisky přinesly element-wise a normalizační kernely. Autoři také uvádějí, že spotřeba tokenů při konverzi klesla v průměru zhruba na polovinu. Každý operátor byl podle nich navíc ověřen z hlediska numerické správnosti.
Omezení, která autoři uvádějí
Autoři sami upozorňují, že ne všechny dodané kernely v TileGym používají plně bezpečný styl. Port totiž musí přesně reprodukovat referenční Tile IR. Tam, kde to dokáže jen nekontrolované API, se proto použije právě ono. Tyto kernely se zatím migrují na bezpečný povrch popsaný v blogu. Za hranicí C ABI navíc podle autorů chybí záchranná síť. Špatný krok, například chybně zadaný stride, proto znamená tichou korupci dat, nikoli výjimku.
Co z toho plyne
Uvedené výsledky podávají samotní autoři a týkají se 24 operátorů z TileGym měřených na DGX B200. Podle autorů překlad nezačíná od nových optimalizací, ale od společného mezijazyka. Jeho výstup lze navíc porovnat strukturálně, nejen testy. Otevřenou položkou zůstává dokončení migrace kernelů, které dnes ještě používají nekontrolované API, na bezpečný povrch.
Grafy a schéma pocházejí z původního článku uvedeného ve zdrojích.
Zdroje
- Translating CUDA Tile Operations from Python to Rust Using Agentic AI | NVIDIA Technical Blog — NVIDIA Developer Blog
Článek vznikl s pomocí AI; podklady a odkazy jsou uvedeny výše.