MoonMath AI új gyorsítókörrel verte az AMD saját figyelmi-magját MI300X-en

A MoonMath AI egy új, nyílt forráskódú HIP figyelmi-magot adott ki az AMD MI300X GPU-khoz. A fejlesztők MIT-licenc alatt használhatják a kódot, amely minden tesztelt alakzatban és kerekítési módban gyorsabbnak bizonyult az AMD saját AITER v3 magjánál. A sebességnövekedés átlagosan 1,18-szoros, de elérheti a 1,26-szorost is. (MarkTechPost)

A mag a transformer modellekben kulcsfontosságú figyelmi műveletet (softmax(QKᵀ/√d)·V) végzi bf16 precizitáson. A MI300X CDNA3 architektúrájára optimalizálták, és a „one-instruction asm wrappers” technikát használja, amely lehetővé teszi az opcode kiválasztását, miközben a fordító kezeli a regisztereket. Ez a megközelítés jelentős sebbesség-növekedést eredményezett a memóriaelhelyezés optimalizálásával: a K bemenet a LDS-ben, a V a L1 gyorsítótárban, a Q és az akkumulátorok pedig regiszterekben helyezkednek el.

Kapcsolódó: Moonshot AI FlashKDA gyorsulás

Optimalizálási részletek és memóriakezelés

A MoonMath AI közlése szerint a mag numerikailag megegyezik az AITER v3 viselkedésével, beleértve a NaN és Inf kezelést, valamint a determinisztikus kimenetet. Egy valós SGLang PR-ban a mag használatával a Wan2.1 videó diffúziós modell betanítási sebessége 1,23-szorosára nőtt minőségromlás nélkül.

Kapcsolódó: FlashQLA gyorsítás Hopper-en

A kernel telepítése egyszerű, pip segítségével végezhető el, és egy kis API-t biztosít. A mag az indítópályán (stream) fut, így könnyen integrálható nagyobb feldolgozási folyamatokba. A BSHD elrendezés közvetlenül alkalmas diffúziós tenzorokhoz, és a cross-attention bármilyen KV hosszal működik.

Kapcsolódó: NVIDIA Star Elastic változatok

Tesztelés és teljesítményösszehasonlítás

A tesztek MI300X-en, bf16 precizitással és 128-as fejmérettel zajlottak. A MoonMath AI magja átlagosan 1,18-szoros (RTNE), 1,15-szörös (RTNA) és 1,08-szoros (RTZ) gyorsulást ért el az AITER-hez képest. A moduláris MAX-hoz képest az átlagos sebességnövekedés 1,44-szerestől 1,49-szeresig terjedt.

Kapcsolódó: OpenAI és Mistral új modellek

A fejlesztők egy nyolchullámú (eight-wave) folyamat-t alkalmaztak két csoportban, két barrierrel, hogy maximalizálják a mátrixmag kihasználtságát. Ez a megközelítés, hasonlóan a FlashAttention-3-hoz, biztosítja, hogy a feldolgozó egység folyamatosan dolgozzon, minimalizálva az üresjáratokat.

Kapcsolódó: Perplexity tokenizálási sebesség

A MoonMath AI által kiadott figyelmi-mag elérhető a GitHubon, és a fejlesztők már használhatják az AMD MI300X GPU-kon a transformer modellek sebességének növelésére.