A PyTorch közvetlenül támogatja a FlyDSL-t az AMD MI350 GPU-kon

Opcionális háttér a PyTorch operátorai mögött

A PyTorch 2026. október 21-i közlése szerint a FlyDSL az AMD MI350 sorozatú GPU-khoz, ezen belül a gfx950 architektúrához kínálható opcionális háttérként. A fejlesztőknek ehhez nem kell alkalmazásaikból egyedi GPU-kerneleket meghívniuk, a támogatott műveletek a meglévő PyTorch API-kon keresztül érhetők el.

A dense GEMM, a grouped GEMM, valamint az MXFP8 és MXFP4 skálázott GEMM a torch.compile használatakor kapcsolódhat a FlyDSL-hez. Az RMSNorm és a TopK megfelelő bemenetek esetén eager módban, automatikus küldéssel használhatja a hátteret. Ha egy bemenet vagy munkaterhelés nem támogatott, a PyTorch megtartja a meglévő megvalósítást.

A TorchInductor a támogatott GEMM-jelölteket kiszűri, ahol lehet, a FlyDSL, az ATen és a Triton teljesítményét is megméri, majd az adott munkaterheléshez a mért eredmény alapján választ és gyorsítótárazza a nyertest.

Az attention támogatása az AOTritonon keresztül érkezik

A FlyDSL attention-kernelei az AOTriton újabb háttérjeként jelennek meg. Ez az a könyvtár, amely a PyTorch ROCm-es F.scaled_dot_product_attention művelete mögött áll, és a FlyDSL-kerneleket az AOTriton 0.14b sorozata tartalmazza.

Az AOTriton az adott munkaterheléshez a FlyDSL, a Triton és az AITER ASM kernelei közül választ az offline operátorhangolási adatbázis alapján. A támogatás az FP16 és BF16 bemenetekre, valamint a 32 és 512 közötti fejméretekre terjed ki. Az FP32 jelenleg nem támogatott, és azok a munkaterhelések, amelyeket a hangolási adatbázis másik megvalósításhoz rendel, továbbra is Triton vagy AITER ASM kernelt használnak.

A PyTorch közlése szerint az AOTriton 0.14.2b gfx950-re közzétett teljesítménye az előre irányú futtatásnál a 0.13b értékének 0,92 és 2,34-szerese, a visszafelé irányú futtatásnál pedig 0,93 és 2,47-szerese között alakult a fejmérettől és a kauzális módtól függően. Az AOTriton FlyDSL attention-kerneleket RDNA4, vagyis gfx1201 GPU-khoz is szállít.

A FlyDSL közvetlenebb hardvervezérlést kínál

A PyTorch magyarázata szerint a Triton blokk- és tenzoralapú programozási modellje a program szintjén megadott csempékre épít, miközben a fordító kezeli a szálak, wave-ok és utasítások jelentős részének leképezését. A FlyDSL ezzel szemben a layout algebrán keresztül teszi explicitté ezeket a leképezéseket.

A kernel szerzője meghatározhatja, hogyan legyenek a csempék blokkok, wave-ok, szálak és értékek között felosztva, kiválaszthatja a tiled-copy és MFMA atomokat, megadhatja az LDS átrendezéseit és az adatbetöltési szakaszokat, valamint a szinkronizációs és ütemezési határokat. A PyTorch szerint ez különösen az architektúraspecifikus sablonoknál hasznos, amikor az adatmozgatás és az utasítások tulajdonjoga az algoritmus részét képezi.

Az attention-kernelek például két wave-os szoftveres csővezetéket használnak. Ez lehetővé teszi, hogy az egyik KV-csempe softmax VALU-munkája átfedje egy másik csempe MFMA-műveletét. A nagyobb fejméreteknél a fejdimenziót az LDS-en keresztül időben, a kimeneti akkumulátort pedig a wave-ok között is felosztják. Ezeket az ütemezéseket közvetlenül a kernelben írják le, így architektúránként hangolhatók a fordító módosítása nélkül.

Mért gyorsulás több művelettípusnál

A PyTorch a kernel szintjén, AMD Instinct MI355X, gfx950 GPU-kon végzett méréseket közölt. A dense GEMM 15 alakzatból álló, BF16-os tesztkészletében a FlyDSL geometriai átlagban 1,10-szer volt gyorsabb az ATen és a Triton alakzatonkénti gyorsabb eredményénél. Az MXFP8 esetében 17 alakzaton 1,58-szoros, az MXFP4 esetében 1,68-szoros geometriai átlagos gyorsulást mértek az ATenhez képest.

A grouped GEMM 24 esetén a Tritonhoz viszonyított geometriai átlagos gyorsulás 1,20-szoros volt. Az RMSNorm 22 esetében az ATenhez képest 1,17 és 3,66 közötti értékeket közöltek. A kis K értékű TopK 10 tesztesetében, nem determinisztikus algoritmusok mellett, a geometriai átlag 4,82-szeres volt az ATenhez képest.

A mérések steady-state operátorfuttatást vizsgáltak, így nem tartalmazták a fordítás és az első automatikus hangolás idejét. Az MXFP eredményeknél a bemenetek már kvantáltak voltak, ezért a kvantálás ideje sem szerepelt a mérési szakaszban. A PyTorch külön kiemeli, hogy ezek az eredmények operátorszintű adatok, és nem azonosak a teljes modell végpontok közötti késleltetésével.

A vLLM végpontok közötti A/B tesztjeiben a teljes kérésre mért gyorsulás BF16-nál elérte az 1,13-szorost, MXFP8-nál a 2,04-szerest, MXFP4-nél pedig a 3,01-szerest. Az eredmény a modelltől és a párhuzamos kérések számától függött.

Mit kapnak a fejlesztők?

A dense GEMM a torch.mm API-n keresztül FP16 és BF16 bemenetekkel, statikus kétdimenziós alakzatokkal, valamint NN, NT, TN és TT elrendezésekkel működik. Az MXFP skálázott GEMM az F.scaled_mm API-t használja, az FP8 E4M3 és a csomagolt FP4 E2M1 formátumokat támogatja, FP32 akkumulációval, és FP16 vagy BF16 kimenetet ad.

A grouped GEMM az F.grouped_mm API-n keresztül eltérő elemszámú csoportokat is kezelhet, így az egyes szakértőkhöz rendelt tokenek száma nem kell, hogy azonos legyen. Az RMSNorm az F.rms_norm, a TopK pedig a torch.topk API-n keresztül érhető el eager módban.

A támogatás minden művelettípusnál alakzat- és igazítási feltételekhez kötött. A nem kompatibilis eager hívások az ATen megvalósításán maradnak, a fordított GEMM-munkaterheléseknél pedig a PyTorch más engedélyezett háttérprogramokat használhat. A FlyDSL így a PyTorch meglévő végrehajtási útvonalaihoz kapcsolódó, mérhető alternatívát ad az AMD MI350 GPU-khoz, miközben a nem támogatott esetek működése változatlan marad.