Šta je FlashAttention sa stanovišta kernel programiranja?

FlashAttention je optimizovani CUDA kernel za izračun mehanizma pažnje (attention) u transformatorima, koji značajno ubrzava izračunavanja i smanjuje potrošnju memorije.

Sa stanovišta programiranja kernela, FlashAttention implementira izračun pažnje koristeći efikasno upravljanje memorijom i paralelizam na GPU. Glavne karakteristike:

• Korišćenje blokova zajedničke memorije za skladištenje međurezultata, što smanjuje pristupe spori globalnoj memoriji.
• Izračun pažnje u nitima sa minimalnim duplikacijama podataka.
• Optimizacija izračuna softmax uzimajući u obzir numeričku stabilnost.
• Obrada velikih sekvenci bez potrebe da se svi podaci učitaju u memoriju odjednom.

Ovo omogućava bržu realizaciju pažnje i manju potrošnju memorije u poređenju sa klasičnim implementacijama, što je ključno za obuku velikih modela.

Na primer, kernel učitava blokove ključeva i zahteva, računa skalarne proizvode, primenjuje softmax i množi sa vrednostima, koristeći zajedničku memoriju za ubrzanje i minimiziranje pristupa globalnoj memoriji.