Co FlashAttention zmienia w inferencji transformerów
Tłumaczenie automatyczne
Ten artykuł został automatycznie przetłumaczony z angielskiego oryginału.
FlashAttention oblicza dokładną uwagę bez zapisywania pełnych macierzy wyników i prawdopodobieństw uwagi w głównej pamięci GPU. Przetwarza mniejsze bloki i łączy ich wyniki za pomocą softmaxu obliczanego na bieżąco. Zmniejsza to ruch w pamięci i ilość przechowywanych danych pośrednich, ale zachowuje kwadratowy koszt obliczeń gęstej uwagi względem długości sekwencji.
Mniejsze zużycie pamięci nie sprawia, że koszt obliczeń gęstej uwagi rośnie liniowo.
Oddziel pamięć od obliczeń
Standardowa materializowana uwaga oblicza wyniki dla każdej widocznej pary zapytanie/klucz, normalizuje je i mnoży przez wartości. Kwadratowa macierz uwagi bez maski zawiera N² elementów dla N tokenów.
Pierwsza publikacja o FlashAttention dzieli te obliczenia na bloki, które mieszczą się w pamięci na chipie. Przechowuje na bieżąco aktualizowane statystyki normalizacji i akumulatory wyjścia zamiast zapisywać wszystkie prawdopodobieństwa w HBM. Przy stałych wymiarach głów pamięć potrzebna na wyniki pośrednie uwagi rośnie liniowo z długością sekwencji. Algorytm nadal oblicza gęste interakcje zapytanie/klucz wymagane przez maskę uwagi.
Algorytm oblicza dokładną uwagę; różna kolejność wykonywania operacji zmiennoprzecinkowych może powodować niewielkie różnice numeryczne. Nie jest to przybliżona metoda rzadkiej uwagi, która pomija wybrane pary tokenów.
Implementacja i obciążenie nadal mają znaczenie
| Pytanie | Dlaczego warto to sprawdzić |
|---|---|
| Czy backend obsługuje ten GPU i tę precyzję? | Implementacje kerneli są przeznaczone dla określonego sprzętu |
| Czy wymiar głów, maska i układ pamięci podręcznej są obsługiwane? | Nieobsługiwane kształty mogą spowodować wybór innego backendu |
| Czy uwaga stanowi dużą część czasu obsługi żądania? | Duże przyspieszenie kernela może mieć niewielki wpływ na całość |
| Czy obciążenie dotyczy fazy prefill, czy dekodowania pojedynczego tokena? | Różni się ilość pracy, którą można wykonać równolegle |
FlashAttention-2 dodatkowo zrównolegla bloki sekwencji zapytań, obok partii i głów, oraz zmniejsza inne narzuty wykonania. FlashAttention-3 wykorzystuje funkcje wykonania charakterystyczne dla architektury Hopper. Nazwy wersji opisują algorytmy i implementacje; nie określają, który backend faktycznie wybrał serwer.
Sprawdź wybrany backend w środowisku serwowania modeli z przypiętą wersją. Zmierz opóźnienie przetwarzania promptu, szczytowe zużycie pamięci i przepustowość całego systemu dla reprezentatywnych długości i rozmiarów partii. Porównuj identyczne maski i precyzję. Benchmark podający wyłącznie TFLOPS kernela nie dowodzi poprawy opóźnienia tokenów odczuwanej przez użytkownika.
Przewodnik inżynierski: FlashAttention omawia kolejne wersje i pomiary w ich określonych zakresach.