Agent alapú cuTile Python és Triton-TileIR kernelok automatizált portolása cuTile Rustra

A TileGym csapata egy több-ügynökös (multi-agent) „skill”-t fejlesztett, amely automatikusan lefordítja a TileGymben megtalálható cuTile Python és Triton-TileIR (nvtriton) tile alapú CUDA kernel implementációkat cuTile Rust-ra (cutile-rs). A konverziós folyamat 24 nyilvános TileGym operátort portolt Rustba (összesen kb. 40 GPU kernel, több kernel esetenként külön variánsokkal), és a mért CUPTI device-idő alapján a cuTile Python referencia átlagos teljesítményének 99,5%-át (geometriai átlag: 0.995) érte el.

Mi az a cuTile Rust és miért fontos a portolás?

cuTile Rust (cutile-rs) egy tile-alapú rendszer, amely idiomatikus, biztonságos GPU-kernel írást tesz lehetővé Rustban. Kiterjeszti a Rust ownership modelljét a tile-alapú GPU-kernelre: a változtatható outputokat diszjunkt darabokra bontja, és fenntartja a host-oldali ownership szerződést kernelindítások között. A felület lehetővé teszi a szerződés alóli lokális kiszállást is alacsony szintű vezérléshez (például közvetlen Tile IR műveletek futtatásához).

A TileGym kódállományában sok produkciós kernel található cuTile Pythonban és Triton-TileIR-ben; a cél az volt, hogy ezek elérhetők legyenek Rustban is, ugyanazzal a teljesítménymodelllel és verifikálhatóan megegyező Tile IR-lettel.

A három front-end és a közös Tile IR

cuTile Python, Triton-TileIR és cuTile Rust három különböző front-end ugyanarra a belső reprezentációra (CUDA Tile IR, a cuda_tile dialektus) épülnek. Mindhárom front-end a tileiras fordítót táplálja, amely végrehajtja a tile-szintű optimalizációkat és előállítja a GPU binárist (cubin). Ez az osztott alap teszi gyakorlatilag és ellenőrizhetővé a front-endek közötti fordítást.

Mivel ugyanazt az IR-t kell reprodukálni, a portolás elsősorban a kernel újrafogalmazása egy biztonságosabb host nyelvben (Rust), nem pedig újra-optimalizálás. Egy hű port célja, hogy a referencia kernel Tile IR-struktúráját reprodukálja: ugyanazok a memória-művelet-családok, tile-alakok és redukciók. A referencia és a fordított kernel Tile IR-ét le tudják dumpolni és egymással diffelni, ez a lépés kulcsfontosságú a pipeline-ban.

A legnagyobb nyelvi szakadék: implicit JIT vs. explicit AOT

A fő különbség az, hogy cuTile Python JIT-ként implicit módon specializál a híváskori paraméterekre, míg cuTile Rust-ban minden specializációt a kernel szignatúrájában (const generics, típusparaméterek) kell deklarálni. Ennek következményei a portolás során:

• Branch-eket, konstansokat és tile-rank-releváns különbségeket a Rustban explicit variánsokra kell felbontani.
• A dinamikus dtype-támogatás helyett a Rust FFI dispatch egy rögzített szimbólum/dtype táblán keresztül dolgozik; új dtype támogatása explicit ABI-kiterjesztés.
• A JIT típusellenőrzése helyett a Rust oldalán vannak két rétegű védelmek: wrapper szintű szemantikai ellenőrzések és FFI-oldali ABI ellenőrzések (null/dtype/device visszakódolásokkal).

A cikk egy konkrét, egyszerű softmax-kernel példán mutatja be a megfeleltetéseket: konstansokból const-generic, ct.load padding_mode explicit partition view + load, ct.astype-ból convert_tile, ct.bid(0)-ból get_tile_block_id(), és a keepdims redukció explicit reduce+reshape+broadcast lépésekké alakul.

C-ABI és integráció a TileGym-mal

A Rust kernelek önállóan is használhatók a cutile crate-nek megfelelő, típusos API-val, de a C-ABI réteg szükséges a TileGym Python dispatch és tesztinfrastruktúrájához. Minden operátor egy C-szimbólumot exportál az aggregált cdylib-ből (például libcutile_kernels.so). A Tensor-ek egyszerű leíró struktúraként (ptr, ndim, shape[], strides[]) haladnak át a határon.

A Python oldalról a cffi egy cdef stringen keresztül köti be a szimbólumot; a Python wrapper nem más, mint vékony validáló réteg, amely ellenőrzi a desc-eket, a device-ot, és elindítja a Rust C-ABI függvényt anélkül, hogy másolatot vagy tulajdonjogot venne. A kernel aszinkron a hívó CUDA stream-jén fut, és a Rust oldalon a borrow_f32 úgy csomagolja a PyTorch eszközmutatót, hogy a Rust ne szabadítsa fel a memóriát, amit a Python birtokol.

A fejlesztési ciklus is kényelmes: cuTile Rust források lusta kompilálással működnek a TileGymben, így a következő hívás automatikusan újrafordítja a shared library-t, ha a forrás megváltozott.

Hogyan működik az agent skill? (felépítés és szerepek)

A tilegym-converting-python-to-rust skill azt a tervezési döntést követi, hogy a top-level agent nem végez mérnöki munkát: csak vezetőként (orchestrator) működik, és specializált alügynököket (subagents) spawnol, amelyek mindegyike csak az adott szakasz referenciáit és utasításait tölti be.

Fő alügynök típusok és feladataik:

• Analyzer: megoldja a „JIT elrejti a specifikációt” problémát. Dumpolja a referencia Tile IR-ét, benchmarkolja a cuTile Python vs Triton-TileIR referencia implementációkat, és kiválasztja leggyorsabb baseline-t variánsonként. Készít egy analysis.json fájlt, ami a variánsokat, konstansokat, dtype-okat, toleranciákat, launch grid-eket és autotune teret rögzíti.
• Kernel writer: előállítja a kernel.rs fájlt. Csak a device-kódra kell koncentrálnia; a translation gap-ot a skill 49 kódolási szabályára alapozva kezeli. Kétféle önellenőrzést futtat: egy in-Rust pipeline tesztet (nincs FFI/Python), és egy struktúrális IR-önellenőrzést a referencia Tile IR ellen.
• Host/FFI builder: elkészíti a C-ABI launchert és a Python wrapper-t, futtatja az operátor valós TileGym tesztkészletét az összes dtype és shape kombináción, és csak az ALL_PASS esetén engedélyezi a benchmarkolást.
• Performance validator: CUPTI protokollt futtatva méri az eszköz-időt páronként ugyanazon GPU-n, és megköveteli, hogy a geometriai átlag 5%-on belül legyen a referenciához képest.

Két diagnosztikai specialista csak hibák esetén kapcsolódik be (nem írnak kódot):

• IR-diff analyst: amikor a correctness teszt vagy a benchmark megbukik, összeveti a referencia Tile IR-t és a generált IR-t, kategorizálja az eltéréseket, és megmondja, hogy kernel-hiba vagy upstream fordítóhiba áll-e a háttérben.
• Residual-performance investigator: ha a kernel helyes, de lassú bizonyos inputokon, arra fókuszál, hogy eszköz-oldali (memória-op család, codegen) és host-oldali (launch, autotune, wrapper) okokra bontva jelentést adjon.

Tervezési elemek, amelyek biztosítják a megbízhatóságot

• Alügynökök csak fix sémájú fájlokon keresztül kommunikálnak, soha nem szabadbeszélgetéssel.
• Minden szakasz egy géppel ellenőrizhető verdict-dal zárul (<VALIDATOR_OUTPUT> blokk és VERDICT: sor); a vezérlő csak erre a verdictra támaszkodik.
• Az osztott cuda_tile dialektus lehetővé teszi, hogy a Tile IR diff legyen a verifikáció gerince: eleinte a kernel writer saját IR self-check-jeként, majd az IR-diff analyst mély összehasonlításaként.
• Bounded retry: minden szakaszra maximált újrapróbálási szám van, így a folyamat vagy konvergál, vagy diagnózissal leáll.

Az orchestrátor loop röviden

1. Preflight ellenőrzés a környezeti változókról és toolchainről; hiba esetén leállás.
2. Minimális pointerekkel spawnolt alügynökök: minden alügynök csak a saját stage-0 fájllistáját és előző szakasz artifactjait kapja meg.
3. Mechanikus validator: az alügynökök returnje egy VALDIATOR_OUTPUT blokkot és egy VERDICT sort tartalmaz; hibás formátum esetén egy automatikus javítási újraindítást kap az alügynök.
4. Routing a döntési táblázat alapján; a failed-perf esetén egy diagnosztikus alág fut.
5. A futás végén a validate_kernel.sh újravizsgálja a teljes 17 fájlos kimeneti szerződést (reportok, IR dumpok, correctness és perf logok).

A repository hierarchiája a skill körül strukturált: SKILL.md, stage instrukciók, referenciák (coding-rules, op-mapping, ir-diff-checklist), koncepciók, script-ek és példák (softmax, bmm) találhatók a skills/tilegym-converting-python-to-rust/ útvonalon.

Példa: softmax kernel megfeleltetés

A cikk egy egyszerű softmax példa bemutatásával szemlélteti a megfeleltetéseket: a cuTile Python dekorált @ct.kernel verziót egy cutile::entry attribútummal ellátott Rust függvényre ültetik át. A különbségek jól láthatók: konstans paraméterek const generics-szé válnak, a padding_mode=NEG_INF explicit partition view + load lépéssé alakul, a típuskonverziók convert_tile-lel történnek, redukció keepdims viselkedése explicit reduce+reshape+broadcast sorozat.

A lényeg: mivel mindkét front-end vékony felszín a cuda_tile IR operátorok felett, a mapping tiszta és a generált Rust IR ugyanazokat az op-kat kell hogy tartalmazza, amit egy IR diff ellenőriz.

Eredmények és benchmarkok

A pipeline CI benchmarkjai CUPTI device-időt mértek NVIDIA DGX B200 rendszeren egy GPU kizárólagos használatával. 347 páros konfigurációt futtattak le a 24 operátoron (minden konfigurációból a CI négy futást vesz, és a legjobb mérést tartják meg). Az összes operátor geometriai átlaga 0.995 volt (paritás a cuTile Python referenciával). Minden operátor áthaladt a 0.95 teljesítményküszöbön; mintegy egyharmaduk gyorsabban futott Rustban, különösen az elem-szintű és normalizációs kernelnél voltak nyereségek.

Fontos megjegyzés: a jelentés CUPTI device-időt mutat, ami a mag kernel futási idejét izolálja. A wall-clock idő más kérdés (tartalmazhat indítási és ütemezési költségeket) — a CI célja az, hogy operátoronként a kernel eszköz-oldali teljesítményét hasonlítsa össze.

Hogyan kezdjünk hozzá?

A skill és az összes konvertált operátor megtalálható a TileGym repo-ban: skills/tilegym-converting-python-to-rust/. A kódok a src/tilegym/ops/cutile_rs/ alatt érhetők el, minden operátornak van egy <op>_kernel/ mappája, és létezik egy aggregált cutile_kernels crate.

Minimális követelmények a futtatáshoz: CUDA 13.1+, Blackwell architektúrájú GPU a perf ellenőrzéshez, Rust 1.89+, és a tileiras compiler.

Induláshoz egy ügynöknek csak annyit kell mondani a repo-ra mutatva, hogy “add a cutile-rs backend for <op>”, a pipeline elvégzi az analízist, kernel-írást, FFI/host építést, validálást és benchmarkolást. A README a TileGym repóban tartalmaz további részleteket.

Záró megjegyzések

A megközelítés legfőbb erénye a megismételhetőség és a felelősségkiosztás: a fordítás minden lépése géppel ellenőrizhető artefaktumokon alapul, és a cuda_tile közös IR lehetővé teszi a strukturális ellenőrzést (IR-diff) a funkcionális tesztek mellett. Ennek eredményeként 24 unattended konverzió futott le sikeresen, a portolt Rust kernellek pedig gyakorlatilag megegyező eszköz-oldali teljesítményt produkálnak a Python referenciával.