Co je FlashAttention z pohledu kernelového programování?

FlashAttention je optimalizovaný CUDA kernel pro výpočet mechanismu pozornosti (attention) v transformerech, který výrazně urychluje výpočty a snižuje spotřebu paměti.

Z pohledu programování kernelů implementuje FlashAttention výpočet pozornosti s využitím efektivního řízení paměti a paralelismu na GPU. Hlavní vlastnosti:

• Použití bloků sdílené paměti pro ukládání mezivýsledků, což snižuje přístupy k pomalé globální paměti.
• Výpočet pozornosti v thready s minimálním duplikováním dat.
• Optimalizace výpočtu softmax s ohledem na numerickou stabilitu.
• Zpracování velkých sekvencí bez nutnosti načítat všechna data do paměti najednou.

To umožňuje provádět pozornost rychleji a s menší spotřebou paměti ve srovnání s klasickými implementacemi, což je kritické pro trénink velkých modelů.

Například kernel načítá bloky klíčů a dotazů, počítá skalární součiny, aplikuje softmax a násobí hodnotami, přičemž využívá sdílenou paměť ke zrychlení a minimalizaci globálních přístupů.