A Together AI 22 petaflops fölé gyorsította GEMM-kernelét Vera Rubinen

A Blackwell örökségére épít a Vera Rubin

A Together AI kernels csapata hozzáférést kapott az NVIDIA Vera Rubin NVL72 platformhoz, majd megvizsgálta az új utasításkészletet és a chip működését. A vállalat szeptember 10-én közzétett beszámolója szerint a ThunderKittensbe olyan funkciókat építettek be, amelyekkel Vera Rubinen NVFP4 és FP8 GEMM-műveletek írhatók.

A Vera Rubin megőrzi a Blackwell programozási modelljét, ezért a korábbi GEMM-kódok is futtathatók rajta. A változtatás nélkül futtatott régi kernel azonban az NVFP4 esetében a platform elméleti teljesítményének körülbelül 42,1 százalékát, FP8 módban pedig 44,4 százalékát érte el. A Together AI szerint ez jelentős teret hagyott az optimalizálásra.

A cég korábbi Blackwell-kernelje klaszterekben futtatja a munkát, az adatokat TMA-n keresztül tölti a megosztott memóriába, a tenzormemóriát használja az akkumulátorok tárolására, és átfedésben kezeli a következő csempék betöltését az előző eredmények kiírásával.

Nagyobb számítási és memória-erőforrások

A Together AI a gyártói specifikációkat összevetve több különbséget emel ki a NVIDIA HGX B200 és a Vera Rubin NVL72 között. Az NVFP4 Tensor Core teljesítménye GPU-nként 9 PFLOPS-ról 35 PFLOPS-ra, az FP8-é 4,5 PFLOPS-ról 17,5 PFLOPS-ra nő. Az FP16 és BF16 Tensor Core értéke 2,25 PFLOPS-ról 4 PFLOPS-ra emelkedik.

A memória-sávszélesség 8 TB/s-ról 22 TB/s-ra változik GPU-nként, a SM-ek száma 148-ról 224-re nő. A megadott csúcsteljesítmény 1000 wattról 2300 wattra emelkedik GPU-nként.

Az egyik fontos újítás, hogy a Tensor Core a K dimenzióban egy lépésben 32 helyett 64 bájtot képes feldolgozni. Ugyanannyi ciklus alatt így kétszer több munka fér el egy utasításablakban, de a Together AI tapasztalata szerint ehhez a kernel adatellátását is fel kell gyorsítani.

A tenzormemória 512-ről legfeljebb 576 oszlopra bővül, ami további 32 KiB kapacitást jelent. A megosztott memória túlméretezett módban 228 KiB-ról 328 KiB-ra növelhető. A Vera Rubin emellett a B oldali csempék újrahasznosítását is támogatja a collector bufferben, valamint lehetővé teszi, hogy az A csempék felszabadításáról a rendszer korábban jelezzen.

Az adatmozgatás lett a fő szűk keresztmetszet

A szélesebb, 64 bájtos K-lépés önmagában csak csekély gyorsulást hozott a korábbi Blackwell-kernelhez képest. A Together AI magyarázata szerint a Tensor Core-ok így kétszer gyorsabban fogyasztják az operandusokat, a kernel azonban nem tudta ugyanilyen ütemben ellátni őket adatokkal.

A vállalat ezért két irányban módosította a működést. Kevesebb adatot mozgatott úgy, hogy ugyanazon CTA-páron egy második kimeneti csempét is elhelyezett az M dimenzió mentén. Így a két akkumulátor megoszthatja ugyanazt a B-csempét. A korábbi 1x1-es csempézés helyett a 2x1-es elrendezés csökkenti az operandusforgalmat.

A mélyebb csővezeték az adatbetöltések jobb átfedését szolgálja. A Vera Rubin új képességeivel együtt ezek a változtatások több mint 22 PFLOPS teljesítményhez segítették az NVFP4-kernelt. A Together AI szerint az eredmény versenyképes a cuBLAS és a CuTE DSL megoldásaival.

Mit jelent ez a felhasználóknak?

A bejelentés elsősorban a nagy teljesítményű gépi tanulási műveleteket fejlesztők számára fontos. A ThunderKittens új támogatása lehetővé teszi, hogy a Vera Rubin hardveres képességeit célzottabb GEMM-kódokkal használják ki, miközben a Blackwellhez kialakított programozási modell megmarad.

A Together AI eredménye azt mutatja, hogy az új GPU-platform teljesítményének kiaknázásához nem elég a szélesebb Tensor Core-utasítás használata. A memória, a csempék újrahasznosítása és a betöltések időzítése egyaránt meghatározza, mennyire közel kerülhet egy kernel a hardver elméleti teljesítményéhez.