Zespół Together AI we współpracy z wiodącymi inżynierami oprogramowania zoptymalizował algorytm FlashAttention-4, osiągając rewelacyjne 71% utylizacji rdzeni graficznych na najnowszych akceleratorach NVIDIA Blackwell B200. To przełomowe rozwiązanie eliminuje sprzętowe "wąskie gardło" asymetrycznego skalowania, pozwalając na trenowanie gigantycznych modeli sztucznej inteligencji nawet 2,7 raza szybciej niż dotychczasowe standardy, co drastycznie obniży koszty infrastrukturalne największych graczy na rynku.
Najważniejsze w skrócie
- Skok wydajnościowy: Algorytm osiąga przepustowość na poziomie 1605 TFLOPs/s na pojedynczym układzie B200, rozwiązując problem niedoboru przepustowości pamięci współdzielonej.
- Nowy paradygmat programowania: Dzięki wykorzystaniu języka CuTe-DSL, zredukowano czas kompilacji kodu o 20-30 razy w porównaniu do tradycyjnych szablonów C++.
- Bezpośrednia integracja: Nowy silnik został zaimplementowany jako backend dla FlexAttention w środowisku PyTorch, co pozwala na tworzenie własnych mechanizmów uwagi bezpośrednio w języku Python.
- Przewaga nad konkurencją: Rozwiązanie wyprzedza autorskie biblioteki NVIDIA (cuDNN 9.13) o 1,1x-1,3x w fazie propagacji w przód (forward pass) oraz zapewnia potężne wzrosty dla dużych okien kontekstowych.
Pułapka asymetrycznego skalowania i ściana pamięci
Ewolucja sprzętowa sztucznej inteligencji napotkała na poważną przeszkodę architektoniczną. Przejście z generacji układów NVIDIA H100 na architekturę Blackwell B200 przyniosło potężny, skokowy wręcz wzrost czystej mocy obliczeniowej. Przepustowość rdzeni Tensor dla operacji w formacie BF16 wzrosła z 1 do 2,25 PFLOPs. Jednak w tym samym czasie przepustowość pamięci współdzielonej dla operacji eksponencjalnych pozostała całkowicie bez zmian.
Ta sprzętowa asymetria stworzyła sytuację, której nikt się nie spodziewał. Karta graficzna zyskała gigantyczny "silnik", ale "rury" dostarczające do niego paliwo (dane) pozostały równie wąskie co u poprzednika. Tradycyjne metody optymalizacji mechanizmu uwagi, które świetnie sprawdzały się jeszcze na układach Hopper (gdzie FlashAttention-3 potrafiło wycisnąć 75% wydajności dla FP16), nagle zderzyły się z murem. Układ B200 przy standardowych algorytmach po prostu czekał bezczynnie na dane, marnując swój potencjał. Badacze odkryli, że w fazie propagacji w przód krytycznym wąskim gardłem stały się obliczenia wykładnicze w funkcji softmax, z kolei w propagacji wstecz (backward pass) całkowicie dominował ruch w pamięci współdzielonej.
Inżynieria pod maską: Jak FlashAttention-4 omija ograniczenia
Aby pokonać fizyczne ograniczenia układów Blackwell, inżynierowie z Together AI, we współpracy m.in. z Tri Dao (twórcą oryginalnego FlashAttention), musieli całkowicie przeprojektować sposób, w jaki oprogramowanie komunikuje się z krzemem. Rozwiązanie to opiera się na innowacyjnym, niskopoziomowym podejściu do zarządzania pamięcią i potokami instrukcji (pipelining).
Optymalizacja propagacji w przód (Forward Pass): Tradycyjny sprzętowy softmax okazał się zbyt wolny. Zespół zastosował harmonogram typu "ping-pong", który przetwarza po dwa kafelki zapytań (query tiles) na każdy blok wątków (CTA). Co najbardziej imponujące, całkowicie zrezygnowano z czekania na sprzętowe jednostki specjalne. Zamiast tego zaimplementowano programową emulację funkcji wykładniczych, wykorzystując instrukcje MUFU.EX2 obok standardowych operacji zmiennoprzecinkowych, co skutecznie podwoiło przepustowość tych obliczeń. Dodatkowo wprowadzono warunkowe skalowanie online – algorytm pomija małe korekty, skracając ścieżkę krytyczną.
Rewolucja w propagacji wstecznej (Backward Pass): To tutaj leżał największy problem z pamięcią współdzieloną. FlashAttention-4 wykorzystuje nową funkcję architektury Blackwell – tryb 2-CTA MMA. Pozwala on na podział akumulatorów wyjściowych między pary bloków. Każdy z nich pobiera tylko połowę operandu B, trzymając jedynie własny "plaster" akumulatora. Wynik? Ruch w pamięci współdzielonej został zredukowany o połowę. Globalne atomowe redukcje dla gradientów również spadły o 50%.
PyTorch i CuTe-DSL: Demokratyzacja badań
Z perspektywy badaczy sztucznej inteligencji, najważniejszą nowością nie są jednak same gigaflopsy, ale sposób, w jaki zintegrowano nową technologię. Zespół wdrożył FlashAttention-4 jako dedykowany backend dla biblioteki FlexAttention w środowisku PyTorch.
Wcześniej, aby stworzyć niestandardowy mechanizm [self-attention] i nie stracić na wydajności, badacze musieli pisać skomplikowany, niskopoziomowy kod w CUDA (często używając abstrakcji CUTLASS C++, które potrafiły kompilować się przez absurdalnie długi czas). Teraz, dzięki zastosowaniu języka CuTe-DSL opracowanego przez zespół NVIDII, czas kompilacji spadł o 20-30 razy. Inżynierowie mogą zdefiniować modyfikatory punktacji (score_mod) lub maski (mask_mod) w zaledwie kilku linijkach czystego Pythona. Kompilator JIT (Just-In-Time) w locie tłumaczy to na zoptymalizowane jądra CuTeDSL pod architekturę Blackwell, zachowując wsparcie dla zaawansowanych technik takich jak ALiBi, maskowanie dokumentów czy technika "sliding window".
W bezpośrednim porównaniu (z włączonym harmonogramem Longest Processing Time first - LPT), nowe rozwiązanie detronizuje dotychczasowe kompilacje Triton, oferując przyspieszenie rzędu 1.5x - 3.2x dla fazy forward i do 2.3x dla fazy backward na układach B200.
Dlaczego to ważne?
Osiągnięcie 71% utylizacji na architekturze NVIDIA Blackwell to nie tylko techniczna ciekawostka dla inżynierów piszących w CUDA – to wstrząs o fundamentalnym znaczeniu dla ekonomii sztucznej inteligencji. Przy obecnej wycenie rynkowej NVIDII na poziomie niemal 4,5 biliona dolarów (stan na marzec 2026 r.) i ogromnym popycie na układy z serii B200, koszt infrastruktury obliczeniowej stanowi główną barierę dla rozwoju AI.
Trenowanie zaawansowanego modelu z rodziny [LLM] pochłania miliony dolarów, z czego ogromna część tych kosztów to praca samych akceleratorów. Gdy karta graficzna pracuje na poziomie np. 40% utylizacji, ponieważ sprzętowe jednostki obliczeniowe czekają na dane z pamięci (efekt tzw. memory wall), firmy dosłownie palą pieniądze w piecu bezczynnego krzemu. FlashAttention-4 rozwiązuje ten krytyczny problem na poziomie algorytmicznym, pokazując, że w dobie asymetrycznego sprzętu nie wystarczy po prostu dokupić więcej kart – konieczny jest inteligentny, sprzętowo-programowy "co-design".
Z perspektywy rynkowej technologia ta drastycznie zwiększa zwrot z inwestycji (ROI) w klastry serwerowe. Jednocześnie głęboka integracja z FlexAttention znosi barierę wejścia dla inżynierów. Oznacza to, że laboratoria badawcze będą mogły znacznie taniej i szybciej eksperymentować z nowymi architekturami sieci neuronowych, co bezpośrednio przełoży się na tempo wdrażania modeli nowej generacji na rynek konsumencki i biznesowy. To ostateczny dowód na to, że optymalizacja oprogramowania dogania w końcu ewolucję sprzętową.
Co dalej?
Mimo rewelacyjnych wyników, zespół programistów nadal intensywnie pracuje nad domknięciem luki między środowiskami badawczymi a produkcyjnymi. W najbliższym czasie możemy spodziewać się:
- Pełnej stabilizacji deterministycznej propagacji wstecznej: Tryb deterministyczny (kluczowy dla powtarzalności treningu) osiąga obecnie 85-90% przepustowości wersji niedeterministycznej. Zespół pracuje nad optymalizacją blokad sprzętowych, aby zniwelować ten spadek.
- Rozszerzenia dla dynamicznych skalarów: Planowane jest wprowadzenie wsparcia dla zmiennych parametrów (np.
soft_cap) w modyfikacjach punktacji bez konieczności ponownej rekompilacji jądra za każdym razem, gdy wartość się zmieni. - Lepszego harmonogramowania w architekturach CUDA: Wykorzystanie klastrowej kontroli uruchamiania (Cluster Launch Control) do sprzętowego, dynamicznego balansowania pracą (zamiast dotychczasowych metod programowych), co jeszcze bardziej przyspieszy operacje na zmiennych długościach sekwencji.
Źródło: Together AI, PyTorch Blog, NVIDIA CUTLASS Team, crypto.news
