Jak mechanizm Flash Attention zmniejsza transfer danych w pamięci podczas lokalnego wnioskowania?

Flash Attention zmniejsza ruch danych w pamięci, dzieląc dokładne obliczenia mechanizmu uwagi na kafelki, dzięki czemu wyniki pośrednie mogą być przetwarzane w układzie zamiast wielokrotnie zapisywane do pamięci HBM GPU i z niej odczytywane.

Na lokalnym serwerze AI efekt ten staje się ważniejszy wraz ze wzrostem długości promptu, ponieważ standardowy mechanizm uwagi tworzy coraz większy obszar roboczy pośredni. Flash Attention nie zmienia wiedzy modelu ani nie usuwa pamięci podręcznej KV; reorganizuje obliczenia mechanizmu uwagi tak, aby przesyłać mniej danych przez najwolniejsze poziomy hierarchii pamięci GPU.

Standardowy mechanizm uwagi przesyła duże tensory pośrednie przez HBM

Mechanizm uwagi rozpoczyna się od tensorów zapytań, kluczy i wartości, a następnie musi obliczyć wyniki, znormalizować je i połączyć z wartościami. Prosta implementacja materializuje duże macierze pośrednie w pamięci o dużej przepustowości między tymi etapami.

W oryginalnej pracy o FlashAttention wskazano, że wydajność mechanizmu uwagi jest ograniczona nie tylko przez obliczenia arytmetyczne, lecz także przez odczyty i zapisy między pamięcią HBM a układową pamięcią SRAM. Długie sekwencje zwiększają ten ruch, ponieważ macierz wyników rośnie wraz z liczbą par tokenów.

Domowy GPU może więc spędzać znaczną część czasu na przenoszeniu danych zamiast na wykonywaniu użytecznych obliczeń macierzowych, szczególnie przy rosnącej długości kontekstu.

Flash Attention dzieli problem na zestawy robocze mieszczące się w pamięci układowej

Flash Attention przetwarza bloki zapytań, kluczy i wartości o rozmiarze dopasowanym do szybkiej pamięci układowej. Każdy kafelek wykonuje kilka etapów obliczeń mechanizmu uwagi, zanim dane pośrednie opuszczą układ.

FlashAttention-3 opisuje podstawowe podejście tej rodziny jako strategię kafelkowania, która eliminuje ruch danych pośrednich przez pamięć globalną. Algorytm nadal oblicza dokładny mechanizm uwagi dla obsługiwanej postaci.

Rozmiar kafelka zależy od zasobów sprzętowych i implementacji kernela. Najważniejszym mechanizmem jest lokalność danych: ponowne wykorzystywanie załadowanych wartości, gdy znajdują się blisko jednostek obliczeniowych.

Softmax online pozwala uniknąć przechowywania pełnej macierzy wyników

Zwykły softmax wydaje się wymagać znajomości wszystkich wyników wiersza przed normalizacją, co zmuszałoby implementację do materializowania dużej macierzy. Flash Attention zamiast tego utrzymuje bieżące statystyki normalizacji dla kolejnych kafelków.

Samouczek Triton dotyczący scalonego mechanizmu uwagi pokazuje scalony mechanizm uwagi z blokowym stanem softmax. Dzięki temu częściowe bloki wyników mogą zostać przetworzone i odrzucone bez zapisywania każdej wartości pośredniej z powrotem do HBM.

W tym miejscu redukcja ruchu danych i redukcja pamięci tymczasowej się łączą: algorytm reorganizuje obliczenia tak, aby duży obiekt pośredni nigdy nie musiał istnieć w pamięci globalnej.

Featured

ZimaCube 2 Osobisty domowy NAS w chmurze

ZimaCube2

$799.00 - $2,499.00

Buy Now

Scalanie kerneli utrzymuje większą część łańcucha obliczeń mechanizmu uwagi w jednej ścieżce wykonania

Obliczanie wyników, skalowanie, maskowanie, normalizacja i agregowanie wartości można połączyć w skoordynowanym kernelu zamiast uruchamiać osobne kernele, które wymieniają pełne dane pośrednie przez pamięć globalną.

Interfejs scaled dot-product attention w PyTorch może wybrać scalone implementacje mechanizmu uwagi, jeśli urządzenie, typ danych, kształt i żądane funkcje są obsługiwane.

Rezultatem może być szybsze przetwarzanie promptów i mniejsze szczytowe zużycie pamięci tymczasowej, ale dokładny zysk zależy od tego, czy środowisko wykonawcze rzeczywiście wybierze kernel zgodny z Flash Attention.

Redukcja operacji wejścia-wyjścia nie oznacza usunięcia pamięci podręcznej KV

Podczas autoregresyjnego wnioskowania wcześniej wygenerowane klucze i wartości nadal muszą pozostać dostępne dla kolejnych tokenów, chyba że model lub środowisko wykonawcze korzysta z osobnego mechanizmu usuwania, kompresji, ograniczania okna lub przenoszenia danych.

Wytyczne AMD dotyczące ROCm opisują Flash Attention jako rozwiązanie poprawiające lokalność pamięci między pamięcią układową a HBM. Optymalizacja ta dotyczy sposobu odczytywania i przetwarzania aktywnego stanu przez mechanizm uwagi, a nie tego, czy historyczny stan KV istnieje.

To rozróżnienie zapobiega częstemu błędowi podczas planowania: Flash Attention może zmniejszyć zużycie pamięci roboczej i liczbę operacji wejścia-wyjścia, a mimo to serwer obsługujący długi kontekst może wyczerpać dostępną pojemność, ponieważ przechowywana pamięć podręczna KV stale rośnie.

Obsługa sprzętu i kształtu danych decyduje o użyciu szybkiej ścieżki

Ker­nele Flash są wyspecjalizowane. Wymiar głowy, typ danych, maska mechanizmu uwagi, architektura GPU, wersja środowiska wykonawczego i inne funkcje mogą wymusić użycie innej implementacji.

Artykuł ZimaSpace o pamięci mechanizmu uwagi w porównaniu z parametrami modelu omawia kwestię pojemności. Flash Attention należy osobno mierzyć jako optymalizację ścieżki wykonania.

Wykonuj testy z krótkimi i długimi promptami, potwierdź wybrany backend oraz rejestruj opóźnienie przetwarzania prefiksu, szczytową przydzieloną pamięć, przepustowość pamięci i szybkość generowania tokenów. Przewodnik ZimaSpace po sprzęcie do lokalnego AI przedstawia szerszą bazę wymagań dotyczących pojemności, której optymalizacja kernela nie może zastąpić.

Transformer Engine firmy NVIDIA omawia zoptymalizowane backendy mechanizmu uwagi i wykonanie zależne od sprzętu jako element szerszego stosu wyboru kernela.