Atencja od strony jądra: FlashAttention, softmax online i połączona dekwantyzacja

Atencja to trzy linijki matematyki:

S = Q·Kᵀ / √d        wyniki, jeden wiersz na zapytanie
P = softmax(S)       wiersz po wierszu
O = P·V              ważona suma wartości

Matematyka jest prosta. Pamięć nie. Poniżej to, co z tym naprawdę robią jądra, od wersji z podręcznika do połączonej ścieżki 4-bitowej, na której działa moja maszyna.

Problem N×N

W naiwnej wersji S i P to macierze N×N dla każdej głowicy w każdej warstwie. Przy prefillu 32K tokenów to 32 768² wyników. W FP16 to 2 GiB na głowicę na warstwę, zapisane do pamięci i od razu z niej czytane. Żadne GPU nie ma tyle pamięci, a nawet gdyby miało, ruch w pamięci zdominowałby wszystko inne.

Wyjście jest jedno: nigdy nie budować tej macierzy. Liczy się ją w kafelkach, które żyją w pamięci na chipie, i wyrzuca każdy po użyciu. Przeszkodą jest softmax: potrzebuje maksimum i sumy całego wiersza, zanim znormalizuje którykolwiek jego element.

Softmax online

Softmax odejmuje maksimum wiersza przed potęgowaniem. To nie jest opcjonalne: w FP16 exp(x) przepełnia się powyżej x ≈ 11,1, bo 65 504 to największa wartość FP16. Softmax online trzyma bieżące maksimum m, bieżącą sumę l i bieżący, nieznormalizowany wynik o, i poprawia je przy każdym nowym kafelku wyników:

m_new = max(m, max_j s_j)
l     = exp(m − m_new) · l  +  Σ_j exp(s_j − m_new)
o     = exp(m − m_new) · o  +  Σ_j exp(s_j − m_new) · v_j
m     = m_new
                                         … po ostatnim kafelku:  O = o / l

Gdy pojawia się większe maksimum, wszystko, co zgromadzono do tej pory, jest przeskalowywane przez exp(m − m_new). Wynik to dokładnie softmax, a nie przybliżenie. Po prostu nigdy nie potrzebuje całego wiersza naraz.

FlashAttention: co nigdy nie trafia do pamięci

FlashAttention (Dao i in., 2022) to softmax online zamieniony w jądro:

1. Załaduj kafelek zapytań do pamięci współdzielonej.
2. Przepuszczaj obok niego kolejne kafelki K i V.
3. Dla każdej pary policz wyniki w rejestrach i zaktualizuj m, l oraz o.
4. Zapisz O raz, na końcu.

S i P istnieją tylko w rejestrach i pamięci współdzielonej, kafelek po kafelku. Jądro czyta Q, K i V i zapisuje O, a nic o rozmiarze N×N nigdy nie trafia do VRAM. Arytmetyka się nie zmienia; zmienia się ruch w pamięci, a na GPU to jest cała gra.

Prefill i dekodowanie to różne problemy

Prefill przetwarza cały prompt naraz. Zapytań jest wiele, więc każdy kafelek K i V jest używany dla wielu wierszy. Praca to duże mnożenia macierzy na rdzeniach tensorowych i ogranicza ją moc obliczeniowa.

Dekodowanie ma jedno zapytanie na sekwencję. Każdy klucz z cache jest użyty do jednego iloczynu skalarnego z tym zapytaniem, a każda wartość do jednego mnożenia z dodawaniem. W FP16 to około 1 FLOP na odczytany bajt, wobec ponad stu FLOP-ów na bajt, które utrzyma współczesne GPU. Atencja przy dekodowaniu to nie mnożenie macierzy, tylko strumień po cache KV, i działa z prędkością pamięci.

Obraz zmieniają dwie rzeczy:

• Grouped-query attention. Gdy g głowic zapytań dzieli jedną głowicę KV, jądro ładuje każdy kafelek K/V raz i używa go g razy. To g× więcej pracy na odczytany bajt.
• Flash-decoding. Przy jednym zapytaniu nie ma czego równoleglić po wierszach, więc jądro dzieli zamiast tego sekwencję KV między bloki wątków. Każdy blok liczy częściowe (m, l, o) na swoim fragmencie, a ostatni krok scala je tą samą regułą przeskalowania co powyżej.

Dlaczego skwantyzowany cache potrzebuje połączonego jądra

Przy pełnym kontekście 262K cache mojego modelu ma 17,2 GB w FP16 i 4,4 GB jako 4-bitowy TBQ4 (wyprowadzenie jest w notatce o cache KV). Atencja przy dekodowaniu czyta go w całości w każdym kroku, więc to format wyznacza liczbę bajtów.

Skwantyzowanego cache można użyć na dwa sposoby:

• Osobny przebieg dekwantyzacji. Przeczytaj 4-bitowy cache, zapisz kopię w FP16, a potem pozwól atencji czytać kopię. Dla całego cache to 4,4 + 17,2 + 17,2 ≈ 39 GB ruchu na krok. To więcej niż przy zwykłym trzymaniu cache w FP16. Oszczędzasz pamięć i tracisz szybkość.
• Dekwantyzacja w trakcie ładowania kafelka. Czytaj 4-bitowe bajty prosto do jądra i zamieniaj je na wartości w rejestrach. Ruch to 4,4 GB i nic więcej.

Jądro TBQ4 z upstreamu w forku Indras-Mirror, na którym buduje mój silnik, robi to drugie, z jeszcze jednym trikiem. Cache jest zapisany w dziedzinie obróconej transformacją Hadamarda, a iloczyny skalarne przeżywają obrót. Zapytanie jest więc obracane raz, klucze i wartości są czytane wprost jako centroid × norma, a wynik jest obracany z powrotem raz. Cała dekwantyzacja to dwa odczyty z tablicy na bajt:

// na bajt = 2 wartości KV
const uint8_t byte = __ldg(&blk->qs[b]);
const half lo = __float2half(d_tbq4_centroids[byte & 0xF] * norm);
const half hi = __float2half(d_tbq4_centroids[byte >> 4] * norm);
tile[...] = __halves2half2(lo, hi);

Przy 262K atencja to wtedy około 21% bajtów czytanych w każdym kroku dekodowania, obok 17 GB wag. Z cache w FP16 byłaby to połowa. Notatka o dekodowaniu zamienia to na sufit tokenów na sekundę.

Zasady, których się trzymam

1. Profiluj prefill i dekodowanie osobno. Mają różne wąskie gardła, a średnia ukrywa oba.
2. Przy długim kontekście atencja przy dekodowaniu to problem przepustowości. Bajty na wartość w cache znaczą więcej niż jakakolwiek optymalizacja FLOP-ów.
3. Skwantyzowany cache opłaca się tylko z połączoną dekwantyzacją. Osobny przebieg dekwantyzacji potrafi przesunąć więcej danych niż cache w FP16, który miał zmniejszyć.
4. Sprawdź proporcję GQA. Wyznacza, ile użycia jądro atencji wyciąga z każdego załadowanego bajtu.

Źródła: Milakov i Gimelshein, Online normalizer calculation for softmax (2018). Dao i in., FlashAttention (2022) i FlashAttention-2 (2023). Dao, Haziza, Massa i Sizov, Flash-Decoding for long-context inference (2023). Fragment jądra TBQ4 pochodzi z README forka upstreamu.