PyTorch pridáva FlyDSL na AMD MI350: kde sa backend vyberá a kedy zrýchli existujúce API

PyTorch podľa svojho technického blogu pripája FlyDSL k existujúcim API tak, aby aplikácie nemuseli volať vlastné kernely. Na AMD MI350-series GPU s architektúrou gfx950 sa to týka dense a grouped GEMM, scaled GEMM pre MXFP8 a MXFP4, RMSNorm, TopK a cez AOTriton aj attention v F.scaled_dot_product_attention. Prakticky to znamená, že zmena nespočíva v novom verejnom rozhraní, ale v ďalšom backende, ktorý môže PyTorch použiť pri podporovaných workloadoch. Podľa dokumentácie PyTorch je táto cesta viazaná na preview/nightly buildy a pri časti funkcionality aj na voliteľný runtime flydsl z vetvy 0.3.x, takže nejde o všeobecne opísanú stabilnú cestu pre všetky inštalácie.

Najdôležitejší mechanizmus je, že FlyDSL nevymieňa celé správanie PyTorch, ale pridáva sa ako kandidát do už existujúcich integračných bodov. Pri eager režime sa podľa blogu podporované volania F.rms_norm a torch.topk môžu automaticky dispatchovať na FlyDSL, ak je backend dostupný, povolený a sú splnené podmienky vstupu; inak zostáva pôvodná implementácia ATen. Pri torch.compile TorchInductor najprv vyfiltruje vhodné dense, grouped a MXFP GEMM workloady, následne benchmarkuje FlyDSL vedľa ATen a Triton, víťaza uloží do cache a spustí najrýchlejší nameraný kernel pre dané zaťaženie. Attention ide inou cestou: F.scaled_dot_product_attention volá AOTriton, ktorý podľa tuning databázy vyberá medzi FlyDSL, Triton a AITER ASM. Praktický význam pre infra tím je v tom, že prínos nezávisí iba od jedného prepínača, ale od konkrétneho operátora, tvaru tensorov, layoutu a od toho, či daný model trávi podstatnú časť času práve v podporovaných operáciách.

Pre dense a grouped GEMM príklad v zdroji výslovne pridáva FlyDSL medzi kandidátov cez import torch._inductor.config as config, config.max_autotune_gemm_backends = "ATEN,TRITON,FLYDSL" a config.flydsl_enable_autotuning = True; kompilovaný výpočet používa torch.compile(mode="max-autotune"). Rovnaké nastavenia zdroj používa aj pre podporované MXFP scaled GEMM. Tieto voľby nezaručujú výber FlyDSL: podľa merania môže vyhrať ATen alebo Triton. Ide o nastavenia Inductora citlivé na verziu nightly, nie o univerzálny prepínač pre eager operácie alebo attention cez AOTriton.

Podpora je úzko vymedzená. Blog uvádza cieľ AMD MI350-series s architektúrou gfx950. Dense GEMM a grouped GEMM sú dostupné cez torch.compile, RMSNorm a TopK v eager režime, attention cez AOTriton a MXFP scaled GEMM vyžaduje už kvantizované operandy aj scale tensory. Pri MXFP ceste platí viac obmedzení naraz: používa sa verejné API F.scaled_mm, podporené sú formáty MXFP8 a MXFP4, akumulácia je vo FP32, výstup je FP16 alebo BF16 a redukčný rozmer K musí byť násobkom 128. Navyše Inductor pri tomto type GEMM pred výberom backendu vyžaduje unit stride pozdĺž K pre A aj B, teda v praxi row-major A a column-major B; pri nekompatibilnom vstupe môže vložiť layout copy. Pri attention je výslovne uvedené, že táto vetva nepodporuje FP32, a torch.nn.attention.flex_attention ide vlastnou Triton cestou, nie cez AOTriton. To sú hranice, ktoré rozhodujú, či sa zmena prejaví bez zásahu do modelu, alebo si vynúti úpravu pipeline.

PyTorch v blogu publikuje vlastné merania na MI355X/gfx950. Na úrovni operátorov uvádza geometrický priemer zrýchlenia 1,10× pre BF16 dense GEMM oproti rýchlejšiemu z ATen a Triton pri každom tvare, 1,20× pre grouped GEMM oproti Triton, 1,58× pre MXFP8 a 1,68× pre MXFP4 scaled GEMM oproti ATen, pri RMSNorm približne 1,17× až 3,66× podľa prípadov a pre malé K pri TopK 4,82× oproti ATen s vypnutou deterministickou voľbou. Samotný blog zároveň upozorňuje, že ide o kernel-level benchmarky, nie o end-to-end latenciu modelu, a pri attention porovnanie verzie AOTriton 0.14.2b proti 0.13b nezachytáva izolovane iba vplyv FlyDSL, ale aj ďalšie zmeny v Triton, AITER a tuning databázach. V end-to-end A/B testoch vLLM na jednom MI355X s pevnou konfiguráciou autor uvádza zrýchlenie celej požiadavky až 1,13× pre BF16, 2,04× pre MXFP8 a 3,01× pre MXFP4, pričom výsledok závisel od modelu a konkurencie požiadaviek. Pre čitateľa je teda podstatné, že ide o doložený mechanizmus s konkrétnymi podmienkami, nie o sľub plošného zrýchlenia každého PyTorch workloadu.