Čo je FlashAttention z pohľadu kernelového programovania?

FlashAttention je optimalizovaný CUDA kernel na výpočet mechanizmu pozornosti (attention) v transformeroch, ktorý výrazne zrýchľuje výpočty a znižuje spotrebu pamäti.

Z pohľadu programovania kernelov, FlashAttention implementuje výpočet pozornosti s využitím efektívneho riadenia pamäti a paralelizmu na GPU. Hlavné vlastnosti:

• Použitie blokov zdieľanej pamäti na uloženie medzivýsledkov, čo znižuje prístupy k pomalej globálnej pamäti.
• Výpočet pozornosti v vláknach s minimálnou duplikáciou dát.
• Optimalizácia výpočtu softmax s ohľadom na numerickú stabilitu.
• Spracovanie veľkých sekvencií bez nutnosti načítania všetkých dát naraz do pamäti.

Toto umožňuje rýchlejšie vykonávanie pozornosti a menšiu spotrebu pamäti v porovnaní s klasickými implementáciami, čo je kritické pre tréning veľkých modelov.

Napríklad, kernel načíta bloky kľúčov a požiadaviek, vypočíta skalárne súčiny, aplikuje softmax a násobí hodnotami, pričom využíva zdieľanú pamäť na zrýchlenie a minimalizáciu globálnych prístupov.