A PyTorch Inductorra váltott az IBM Spyre gyorsító fordítója

Az Inductor lett az új alap

A PyTorch 2026. október 13-án közzétett beszámolója szerint a torch-spyre, vagyis az IBM Spyre gyorsítót célzó PyTorch-integráció, a sendnn helyett immár az Inductor kibővített változatát használja a torch.compile fordítási háttérrendszereként. A sendnn teljes egészében külön fejlesztett fordítási stack volt, az Inductor viszont a PyTorch saját, nyílt forráskódú backendje.

A váltás eredményeként létrejött fordító körülbelül 7-szer kisebb az elődjénél. A PyTorch szerint az Inductor nagy része változtatás nélkül használható maradt, a fejlesztési munka pedig főként azokra a pontokra összpontosult, ahol az IBM gyorsítója alapvetően másképp működik, mint egy GPU.

A torch-spyre fordítója két szakaszból áll. Az első az Inductor kiegészítéseivel a torch.compile által előállított gráfot a SuperDSC nevű köztes reprezentációra alakítja. A második, a DeepTools nevű saját fejlesztésű fordító, a SuperDSC-t a Spyre hardverén végrehajtható utasításokká fordítja. A PyTorch bejegyzése az első szakasz, vagyis az Inductor-bővítések működését mutatja be.

A Spyre más adatmozgatást vár el

Az IBM a Spyre-t vállalati következtetési feladatokra tervezett, energiahatékonyságot előtérbe helyező gyorsítóként mutatja be. A PCIe-n csatlakozó eszköz fogyasztása 75 watt, és a leírás szerint nincs szüksége egyedi tápellátásra vagy hűtésre. A gyorsító adatfolyam-alapú felépítésű, 32 kisebb magból áll, amelyeket gyűrűs összeköttetés kapcsol össze. Mindegyik magnak saját, 2 MB-os helyi munkamemóriája van közös gyorsítótár helyett.

A Spyre az adatokat rögzített méretű egységekben, úgynevezett stickekben mozgatja. Egy stick 128 bájt, ami fp16 formátumban 64 értéknek felel meg. Míg egy GPU-s kernel egyedi elemeket címezhet, a Spyre alapvető adatmozgatási és számítási egysége a stick. A fordítónak ezért az indexelést is ehhez kell igazítania.

Az egyes műveletek meghatározzák, hogy a bemeneti és kimeneti tenzorok elemei miként rendezhetők stickekbe. Egy elemenkénti művelet például azt igényli, hogy a bemenetek és a kimenet ugyanazt a stickdimenziót használják. Ha ez nem teljesül, a fordítónak restickify műveletet kell beillesztenie, amely fizikailag újrarendezi az adatokat, és a stickek szerkezetét is átalakítja.

A stickdimenziók kiosztása ezért a teljes számítási gráfot érintő optimalizációs feladat. A torch-spyre által választott tenzorelrendezés a teljesítményre is hatással van. A PyTorch mérése szerint egy, már a gyorsító memóriájában lévő, kötegelt mátrixszorzásnál a bemeneti és kimeneti tenzorok csempézésének módosítása a kernel futási idejét 29 és 69 százalék közötti mértékben csökkentette különböző kötegméretek és alakok mellett.

A fordítónak a memóriát és a magokat is meg kell terveznie

A Spyre magonkénti scratchpad memóriát használ, amely gyors, de korlátozott méretű. Mögötte a globális eszközmemória található, és az adatok mozgatása nem automatikus. A scratchpad és a számítási egységek közötti, valamint az eszközmemóriából a scratchpadba irányuló minden adatmozgáshoz a fordítónak külön utasításokat kell előállítania.

Ha a fordító nem tervezi meg egy érték scratchpadban tartását, azt vissza kell írni az eszközmemóriába, majd a következő használatkor ismét be kell tölteni. A PyTorch szerint a Spyre feladatainál gyakran az adatmozgatás jelenti a szűk keresztmetszetet, ezért a scratchpad előzetes megtervezése a fordító egyik központi feladata.

A hardver harmadik fontos eltérése, hogy a magok statikus adatfolyamú gépként működnek. Nincs olyan, GPU-khoz hasonló futásidejű ütemező, amely menet közben egyensúlyozza a terhelést a számítási egységek között. A munkát ezért már fordításkor fel kell osztani a magok között.

A PyTorch bővítési pontjaira építettek

A torch-spyre több ponton egészíti ki az Inductor feldolgozási folyamatát. Saját műveletfelbontásokat használ a Spyre-re optimalizált primitív műveletsorok előállításához, majd ezekhez alsó szintű leképezéseket regisztrál. A legfontosabb munka a CustomPreSchedulingPasses szakaszban történik, ahol a csempézésről, az elrendezésről, a munkamegosztásról és a scratchpad kiosztásáról születnek döntések.

Ez a rész jelenleg az Inductor GraphLowering.codegen() metódusának monkey patch alapú módosításán keresztül kapcsolódik a rendszerhez. A PyTorch saját, későbbi bővítési pontja, a _pre_fusion_custom_pass, túl későn futna, mert addigra a köztes reprezentáció SchedulerNodes elemekké alakult, és az Inductor már nem támogatja a gráf nagy léptékű átírását.

A PyTorch és a torch-spyre fejlesztői most azon dolgoznak, hogy formalizálják a használt bővítési pontokat. A cél, hogy más, a fő kódfán kívül fejlesztett hardveres backendek is hasonlóan profitálhassanak az Inductor infrastruktúrájából.