FlashAttention: uwaga, która nie zjada pamięci

Znasz to -fa on, które wklepujesz przy odpalaniu llama.cpp prawie z automatu, jak człowiek zapinający pas w aucie — nie zastanawiasz się, po prostu wiesz, że tak trzeba? To jeden z tych patentów, który powtarza się jak refren: „nie wchodzi na kartę? dorzuć flash attention". Wrzucasz, kontekst się rozciąga, VRAM-u ubywa wolniej, i jedziesz dalej, nawet nie pytając, co ta flaga właściwie robi pod maską.

A pod maską siedzi jedna z ładniejszych prac ostatnich lat — FlashAttention, Tri Dao i spółka ze Stanfordu, maj 2022.¹ Ładna nie dlatego, że wymyśliła nowy rodzaj uwagi. Wręcz przeciwnie — wynik matematyczny wychodzi co do liczby taki sam jak w zwykłej uwadze. Cały myk polega na czymś, co brzmi jak herezja dla kogoś, kto liczył złożoność na studiach: żeby przyspieszyć obliczenia, zamiast liczyć mniej, autorzy kazali karcie mniej biegać po pamięci. I to załatwiło sprawę.

Skąd w ogóle ten głód pamięci

Najpierw ustalmy, czego uwaga potrzebuje, bo tu leży cały problem. Uwaga (ta z „Attention Is All You Need", z 2017 — inna praca, inni ludzie, nie mylić) dla sekwencji o długości N tokenów porównuje każdy token z każdym. Każdy z każdym to macierz N × N. Masz 4 tysiące tokenów kontekstu? To 16 milionów liczb w jednej tablicy. Osiem tysięcy? Już 64 miliony. Pamięć rośnie z kwadratem długości — podwajasz kontekst, płacisz czterokrotnie.¹

I teraz najważniejsze, co ludzie przeoczają: nie chodzi o to, że tych mnożeń jest dużo. GPU uwielbia mnożyć, od tego jest. Problem w tym, gdzie ta wielka macierz ląduje. Karta ma dwa rodzaje pamięci. Jest HBM — ta duża, te twoje 8, 12 czy 24 giga, wolna jak poranny ruch na obwodnicy. I jest SRAM — malutka pamięć tuż przy rdzeniach obliczeniowych, błyskawiczna, ale mieści tyle co nic. Zwykła uwaga liczy całą macierz N × N, zapisuje ją w HBM, potem odczytuje z powrotem, żeby policzyć softmax, znowu zapisuje, znowu czyta… i to bieganie tam i z powrotem — nie samo mnożenie — jest tym, co zżera czas i pamięć.

„IO-aware", czyli licz tam, gdzie blisko

Tu wchodzi pomysł, od którego cała praca dostała podtytuł: IO-awareness, świadomość operacji wejścia-wyjścia.¹ Brzmi korporacyjnie, a znaczy rzecz prostą jak drut: przestań optymalizować liczbę mnożeń, zacznij optymalizować liczbę podróży do wolnej pamięci. Bo to one są wąskim gardłem.

Jak to zrobić, skoro macierz N × N z definicji jest ogromna i nie zmieści się w tej malutkiej szybkiej SRAM? No właśnie — nie zmieści się w całości. Ale w kawałkach już tak. I to jest sedno.

• Tiling (kafelkowanie). Zamiast liczyć naraz całą macierz, tniesz Q, K i V na bloki i przerabiasz je kafelek po kafelku. Każdy kafelek jest na tyle mały, że wjeżdża do SRAM, tam się go przemiela, wypluwa wynik i bierze następny. Wielka macierz N × N jako całość nigdy nie powstaje w wolnej pamięci — istnieje tylko przelotnie, we fragmentach, w tej szybkiej pamięci przy rdzeniach.²
• Online softmax. Tu jest haczyk, bo softmax klasycznie potrzebuje całego wiersza naraz — musi znormalizować po wszystkich elementach. FlashAttention liczy go „w locie", z bieżącą normalizacją: przerabia kolejne bloki i na bieżąco koryguje wynik, tak jakby sumował rachunek pozycja po pozycji, poprawiając napiwek za każdym razem, gdy dojdzie nowa pozycja. Efekt końcowy identyczny, a pełnego wiersza nigdy nie trzeba trzymać.²
• Rekomputacja. Przy uczeniu (backward pass) normalnie trzymasz w pamięci wszystkie pośrednie wyniki, żeby policzyć gradienty. FlashAttention mówi: po co je trzymać, skoro taniej je policzyć jeszcze raz z tego, co już mamy w szybkiej pamięci? Dokłada trochę mnożenia, żeby oszczędzić na pamięci i na tym przeklętym bieganiu do HBM. I na tym akurat wychodzi na plus.¹

Zbierz to razem i dostajesz uwagę, w której zużycie pamięci rośnie liniowo z długością sekwencji, O(N), zamiast kwadratowo, O(N²).¹ Ta sama matematyka, ta sama odpowiedź co do ostatniej cyfry — tylko poukładana tak sprytnie, że wielka tablica pośrednia nigdy nie musi się zmaterializować tam, gdzie jest wolno.

Czy to na pewno działało? Liczby z pracy

Bo można sobie opowiadać o eleganckim algorytmie do rana, ale papier bez liczb to wróżenie z fusów. Autorzy zmierzyli i pokazali czarno na białym.¹

I najciekawsze na koniec — nie samo „szybciej", ale „w ogóle się dało". Skoro pamięć przestała rosnąć kwadratowo, można było pchać modele na kontekst, na którym wcześniejsze Transformery po prostu się poddawały. Na Path-X (sekwencja 16 tysięcy) FlashAttention wyciągnął 61,4% dokładności, a na Path-256 (64 tysiące!) — 63,1%.¹ To zadania, na których wcześniej model radził sobie nie lepiej niż rzut monetą. Dłuższy kontekst, który dziś traktujesz jak coś oczywistego, częściowo wziął się właśnie stąd.

Zrób to sam — i nie pomyl tego z KV-cache

Nie musisz nic implementować, cały ten mechanizm masz w jednej fladze. W llama.cpp to -fa, --flash-attn, przyjmuje on/off/auto, domyślnie stoi na auto.³

llama-cli -m qwen2.5-7b-instruct-q4_k_m.gguf -ngl all -c 8192 -fa on -cnv

Jedna rzecz, którą warto tu dorzucić, bo ludzie się na niej wykładają: w llama.cpp flash attention jest twardym wymogiem dla kwantyzacji KV-cache. Chcesz ścisnąć cache do q8_0, żeby odzyskać VRAM przy długim kontekście? Bez -fa on się nie da — kwantyzacja V-cache wymaga flash attention wprost, kod ucina to jednym komunikatem błędu.⁴ Więc te dwie sztuczki chodzą w parze:

llama-cli -m qwen2.5-7b-instruct-q4_k_m.gguf \
  -ngl all -c 8192 -fa on -ctk q8_0 -ctv q8_0 -cnv

I tu ostrzeżenie, żebyś nie wpadł w częstą pułapkę pojęciową. FlashAttention to nie jest KV-cache — to dwie różne rzeczy, które akurat obie oszczędzają pamięć. KV-cache dotyczy tego, że model przechowuje przeliczony kontekst między kolejnymi tokenami generacji (rozgryzaliśmy to bliżej tutaj). FlashAttention dotyczy tego, jak liczysz uwagę w jednym przejściu — kolejności obliczeń, nie przechowywania stanu. Oszczędność bierze się wyłącznie z tilingu i rekomputacji, a nie z jakiejkolwiek utraty precyzji — nikt tu niczego nie kwantyzuje ani nie zaokrągla.

Dwa zastrzeżenia od siebie, żeby cię to potem nie zaskoczyło. Raz — wsparcie nie jest uniwersalne: na backendzie CUDA jest jak znalazł, ale na Vulkanie czy SYCL dla części kart flaga bywała kapryśna, więc nie traktuj jej jak działającej wszędzie i zawsze (więcej o tym, czym różnią się backendy). Dwa — jak zwykle w llama.cpp, resztę flag, które chodzą obok tej, ogarnialiśmy osobno.

Historia nie skończyła się w 2022. Rok później przyszedł FlashAttention-2 — już samego Tri Dao — który poprzestawiał podział pracy między wątkami GPU i wycisnął jeszcze z grubsza dwukrotne przyspieszenie względem jedynki, dobijając do 50–73% teoretycznego maksimum kart A100.⁵ Kod jednego i drugiego (plus wersji trzeciej) siedzi w oficjalnym repo Dao-AILab/flash-attention, na licencji BSD-3.²

Ale najfajniejsze w tej całej historii jest to, jak przewrotnie brzmi jej morał. Przez lata uczono nas, że żeby program był szybszy, trzeba go zmusić do liczenia mniej. A tu przyszli ludzie, którzy kazali karcie policzyć część rzeczy dwa razy — i wyszło szybciej. Bo okazało się, że najdroższa rzecz na GPU to nie mnożenie liczb, tylko chodzenie po nie na drugi koniec pamięci. Trochę jak z tym kolegą, który zamiast nosić zakupy po jednej siatce z auta pod blok, woli obładować się wszystkim naraz i przejść ten dystans raz. Nie jest leniwy. Po prostu wie, że najgorszy jest sam spacer tam i z powrotem.

Przypisy · źródła