Czym jest FlashAttention z punktu widzenia programowania kerneli?

FlashAttention to zoptymalizowany kernel CUDA do obliczania mechanizmu uwagi (attention) w transformatorach, który znacznie przyspiesza obliczenia i zmniejsza zużycie pamięci.

Z punktu widzenia programowania kernelowego, FlashAttention realizuje obliczanie attention z wykorzystaniem efektywnego zarządzania pamięcią i równoległości na GPU. Główne cechy:

• Użycie bloków pamięci współdzielonej do przechowywania danych pośrednich, co zmniejsza odwołania do wolnej pamięci globalnej.
• Obliczanie attention w wątkach z minimalnym duplikowaniem danych.
• Optymalizacja obliczania softmax z uwzględnieniem stabilności numerycznej.
• Obsługa dużych sekwencji bez konieczności ładowania wszystkich danych do pamięci jednocześnie.

Pozwala to na szybsze wykonywanie attention i mniejsze zużycie pamięci w porównaniu z klasycznymi implementacjami, co jest krytyczne dla trenowania dużych modeli.

Przykładowo, kernel ładuje bloki kluczy i zapytań, oblicza iloczyny skalarne, stosuje softmax i mnoży przez wartości, korzystając z pamięci współdzielonej w celu przyspieszenia i minimalizacji odwołań do pamięci globalnej.