FlashAttention-3 (HU)

FlashAttention-3 — egy algoritmus a transzformeres neurális hálózatok attention mechanizmusának optimalizálására, amelyet az NVIDIA Hopper (H100) architektúrájú GPU-k hardveres lehetőségeinek maximális kihasználására fejlesztettek ki. Az algoritmust 2024-ben mutatta be egy kutatócsoport a Colfax Research, a Meta, az NVIDIA, a Georgia Tech, a Princetoni Egyetem és a Together AI vállalataitól. A munkát elfogadták a NeurIPS 2024 konferenciára, ahol spotlight minősítést kapott.

A FlashAttention-3 az algoritmuscsalád harmadik iterációja, amely a FlashAttention (2022) és a FlashAttention-2 (2023) után következik. Fő célja a nagy nyelvi modellek (LLM) tanításának és inferenciájának jelentős gyorsítása, a számítási pontosság megőrzése mellett.

Bevezetés és előzmények

Az attention mechanizmus problémája

A transzformerek kulcskomponense az önfigyelmi (self-attention) mechanizmus, amelynek számítási komplexitása és memóriaigénye azonban négyzetes mértékben (O(n²)) növekszik a bemeneti szekvencia hosszával (n). Ez komoly „szűk keresztmetszetet" jelent, mivel a modern GPU-k gyors mátrixszorzásokra vannak optimalizálva, de az exponenciális függvények kiszámítása (például a Softmax-ban) nagyságrendekkel lassabb. Emellett a naiv implementáció esetén a GPU memóriájában nagy méretű közbülső attention tenzort kell tárolni, ami korlátozza a modellek skálázhatóságát.

FlashAttention és FlashAttention-2

Ennek a problémának a megoldására 2022-ben javasolták a FlashAttention algoritmust, amely két technika segítségével csökkentette a lassú globális memóriához (HBM) való hozzáférések számát:

• Blokkos feldolgozás (tiling): A számítások blokkokra (tile-okra) vannak osztva, amelyeket a gyors on-chip memóriában (SRAM) dolgoznak fel.
• Műveletek összevonása: Az összes művelet (mátrixszorzás, Softmax) egyetlen GPU kernel-ben fut le, anélkül hogy a közbülső eredményeket a globális memóriába írnák.

Ez lehetővé tette, hogy a memória-komplexitás négyzeteséről lineárisra csökkenjen, és a számítások 2–4-szeres gyorsulást értek el.

2023-ban bemutatták a továbbfejlesztett verziót — a FlashAttention-2-t —, amely optimalizálta a számítások párhuzamosítását. Az NVIDIA Ampere (A100) architektúrájú GPU-n az elméleti csúcsteljesítmény ~70%-át érte el. Azonban az újabb NVIDIA Hopper (H100) architektúrán a hatékonysága lényegesen alacsonyabb volt — körülbelül 35%. Ez azzal volt összefüggésben, hogy az algoritmus nem használta ki a Hopper új hardveres lehetőségeit, ami ösztönzőül szolgált a FlashAttention-3 létrehozásához.

A Hopper (H100) GPU új hardveres lehetőségei

Az NVIDIA Hopper architektúra számos új funkciót biztosított, amelyeket a FlashAttention-3 a maximális teljesítmény elérése érdekében kihasznál:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): Új típusú utasítások a tenzormagok számára, amelyek mátrixszorzásokat hajtanak végre, közel kétszeres teljesítménynövekedést nyújtva az Ampere architektúrához képest.
• TMA (Tensor Memory Accelerator): Hardveres modul, amely gyorsítja az adatátvitelt a globális (HBM) és a megosztott (shared memory) memória között. A TMA automatikusan elvégzi a cím-számításokat, tehermentesítve a számítási kerneleket.
• FP8 formátum: 8 bites lebegőpontos adatformátum hardveres támogatása, amely megkétszerezi az elméleti teljesítményt az FP16-hoz képest, de a korlátozott dinamikus tartomány miatt a pontosság elvesztésének kockázatával jár.

A FlashAttention-3 technikai újításai

Az algoritmus három kulcsfontosságú optimalizálási módszert valósít meg, amelyeket kifejezetten a Hopper architektúrára fejlesztettek ki:

1. Aszinkron végrehajtás és warp specializáció

A FlashAttention-3 a warp-specialization elvét alkalmazza, amelynek során a GPU-n lévő különböző szálcsoportok (warps) különböző feladatokra specializálódnak:

• Producer warps: Adatokat töltnek be a globális memóriából TMA segítségével.
• Consumer warps: Mátrixszorzásokat hajtanak végre a tenzormagokon.

A Hopper hardveres aszinkronitásának köszönhetően ezek a műveletek időben átfednek. Amíg az egyik warp-csoport számításokat végez, a másik párhuzamosan tölti be a következő blokk adatait. Ez a csővezeték-szerű megközelítés (pipeline), amelyet „ping-pong ütemezésnek" (ping-pong scheduling) neveznek, lehetővé teszi a lassú műveletek (pl. Softmax) késleltetésének elrejtését, és a GPU összes funkcionális moduljának maximális kihasználását.

2. Memóriaműveletek minimalizálása

Az algoritmus megőrzi a korábbi verziókból ismert tiling szemléletet, de aktívan használja a TMA-t a következő adatblokkok aszinkron betöltésére, párhuzamosan az aktuális számításokkal. A lassú HBM-ből a gyors SRAM-ba történő adatátvitel gyakorlatilag a fő számítások „árnyékában" zajlik, ezáltal a GPU kevesebbet vár az adatokra.

3. Alacsony pontosság (FP8) a kvantálási hiba csökkentésével

Az FP8-ra való áttérés megduplázza a sebességet, de a kvantálás miatt jelentős pontosságvesztéshez vezethet. Ennek kezelésére a fejlesztők bevezették az incoherent processing módszert. Ennek lényege a következő:

1. Az attention kiszámítása előtt a jellemzővektorokat (Q lekérdezések és K kulcsok) megszorozzák egy véletlenszerű ortogonális mátrixszal (például Hadamard-mátrixszal).
2. Ez az átalakítás „szétoszlatja" a rendkívül nagy abszolút értékű elemeket (kiugró értékeket) az összes koordinátára, kiegyenlítve azok eloszlását.
3. Ezt követően megtörténik az FP8-ra való kvantálás, amely most kisebb hibával jár.
4. Mivel az átalakítás ortogonális, nem torzítja az attention végeredményét (QKᵀ), mivel a mátrix hatása megszorzáskor kiegyenlítődik.

Ez a technika az attention FP8-ban való kiszámításának hibáját körülbelül 2,6-szorosára csökkentette az átalakítás nélküli szabványos FP8 alkalmazáshoz képest.

Teljesítmény és jelentőség

A felsorolt technikák alkalmazása lehetővé tette, hogy a FlashAttention-3 jelentős fölényt érjen el a korábbi verziókhoz képest H100 GPU-n:

• 1,5–2-szeres gyorsulás a FlashAttention-2-höz képest.
• Magas GPU-kihasználtság: A H100 elméleti maximális teljesítményének ~75–85%-át éri el.
• Átviteli sebesség:
  • Legfeljebb 740–840 TFLOPS fél pontosság esetén (FP16/BF16).
  • Legfeljebb 1,2–1,3 PFLOPS (petaflops) 8 bites pontosság (FP8) használatakor.

A FlashAttention-3 magas hatékonysága közvetlenül befolyásolja az LLM-ek fejlesztését és alkalmazását:

• Tanítási idő csökkentése: Az attention 75–100%-os gyorsítása jelentősen lerövidíti a modellek tanítási idejét, amely hetekig vagy hónapokig tarthat.
• Kontextusablak növelése: A modellek hatékonyan képesek feldolgozni hosszabb szekvenciákat (több százezer token), ami fontos nagy dokumentumok vagy kód elemzéséhez.
• Erőforrások ésszerű felhasználása: Lehetővé teszi ugyanolyan teljesítmény elérését kevesebb GPU-val, vagy nagyobb sebesség elérését ugyanazon hardveren, ami csökkenti a modellek bevetési költségét.

Elérhetőség és integráció

A szerzők nyílt licenc alatt publikálták a FlashAttention-3 forráskódját a GitHub-on. Várható az integráció a vezető mélytanulási keretrendszerekbe, mint a PyTorch és a Hugging Face Transformers könyvtárai, ami széles körű fejlesztők és kutatók számára elérhetővé teszi a technológiát. A korábbi verziók már de facto szabvánnyá váltak az iparban, és a FlashAttention-3 valószínűleg folytatja ezt a tendenciát.

Hivatkozások

• A FlashAttention hivatalos GitHub repozitóriuma
• A Together AI blogja a FlashAttention-3 bejelentésével

Irodalom

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.