Mi az a FlashAttention a kernel-programozás szempontjából?

A FlashAttention egy optimalizált CUDA kernel a figyelem (attention) mechanizmus számítására a transzformátorokban, amely jelentősen felgyorsítja a számításokat és csökkenti a memóriafelhasználást.

Kernel programozási szempontból a FlashAttention hatékony memóriakezelést és párhuzamosságot alkalmaz a GPU-n a figyelem számításához. Fő jellemzői:

• Megosztott memória blokkok használata ideiglenes adatok tárolására, ami csökkenti a lassú globális memória hozzáféréseket.
• Figyelem számítása minimális adatduplikációval szálakban.
• A softmax számításának optimalizálása a numerikus stabilitás figyelembevételével.
• Nagy szekvenciák kezelése anélkül, hogy az összes adatot egyszerre kellene betölteni a memóriába.

Ez lehetővé teszi a figyelem gyorsabb végrehajtását és kevesebb memória felhasználását a klasszikus megvalósításokhoz képest, ami kritikus a nagy modellek tanításához.

Például a kernel betölti a kulcsok és kérdések blokkjait, kiszámítja a skaláris szorzatokat, alkalmazza a softmaxot és megszorozza az értékekkel, a gyorsítás és a globális memória hozzáférések minimalizálása érdekében megosztott memóriát használva.