Mnożenie macierzy i jak działają akceleratory tensorowe

Odrzuć nazwy, a LLM okaże się głównie jedną operacją, powtarzaną miliardy razy: mnożeniem dwóch macierzy. Projekcje atencji, MLP i głowica wyjściowa to same operacje matmul. Każdy akcelerator zbudowany z myślą o AI jest w gruncie rzeczy maszyną do szybkiego wykonywania tej jednej operacji. Jeśli rozumiesz, ile kosztuje matmul, rozumiesz większość tego, co widzisz w wydajności inferencji i treningu.

Ile kosztuje matmul

Pomnożenie macierzy m × k przez macierz k × n daje m × n wartości wyjściowych. Każda z nich to iloczyn skalarny długości k, czyli k mnożeń i k dodawań:

FLOPs = 2 · m · n · k
bytes ≥ (m·k + k·n + m·n) · bytes_per_element      odczyt A i B, zapis C

Stosunek tych dwóch wielkości to intensywność arytmetyczna, czyli liczba FLOP-ów na przesłany bajt. To ona decyduje o wszystkim:

• Duży kwadratowy matmul, 4096 × 4096 × 4096 w FP16: 137 GFLOP na około 100 MB danych, czyli około 1365 FLOP-ów na bajt. Na RTX 4090 to około 0,83 ms obliczeń wobec 0,1 ms ruchu w pamięci. To mnożenie jest ograniczone mocą obliczeniową.
• Ta sama macierz razy jeden wektor, czyli dekodowanie przy partii o rozmiarze 1: 33,5 MFLOP na 33,5 MB wag, 1 FLOP na bajt. To około 33 µs ruchu w pamięci wobec 0,2 µs obliczeń. To mnożenie jest ograniczone pamięcią, i to mniej więcej 160-krotnie.

Te same wagi, ten sam sprzęt, dwa zupełnie różne wąskie gardła. Na tym polega cała różnica między prefillem a dekodowaniem (konsekwencje opisuje notatka o dekodowaniu).

Model roofline

Model roofline (Williams, Waterman i Patterson, 2009) ujmuje to w jednym obrazie. Osiągalna wydajność to niższy z dwóch sufitów:

osiągalne FLOP/s = min( szczytowe FLOP/s ,  intensywność arytmetyczna × przepustowość pamięci )

Punkt załamania (ridge point), w którym spotykają się oba sufity, to szczyt / przepustowość. Poniżej niego ogranicza Cię przepustowość pamięci, a powyżej moc obliczeniowa.

GPU	Szczytowa wydajność tensorowa (16 bitów, dense)	Przepustowość pamięci	Punkt załamania
RTX 4090	~165 TFLOPS (akumulacja w FP32)	1008 GB/s	~164 FLOP-y na bajt
H100 SXM	~989 TFLOPS (BF16)	3,35 TB/s	~295 FLOP-ów na bajt

Dekodowanie pojedynczego strumienia to około 1–4 FLOP-ów na bajt, o dwa rzędy wielkości poniżej każdego z tych punktów załamania. Dlatego batching podnosi przepustowość prawie za darmo: B sekwencji mnoży intensywność przez B (liczby są w notatce o batchingu). I dlatego nowsze GPU, których moc obliczeniowa rośnie szybciej niż przepustowość pamięci, odsuwają punkt załamania jeszcze dalej poza zasięg dekodowania.

Kafelkowanie: ponownie używaj tego, co już załadowane

Naiwne jądro czyta każdy element A i B z pamięci raz na każdą wartość wyjściową, do której się przyczynia, a to o wiele za dużo ruchu. Prawdziwe jądra dzielą wynik na kafelki. Fragmenty A i B potrzebne danemu kafelkowi są ładowane raz do szybkiej pamięci na chipie (pamięci współdzielonej i rejestrów) i używane ponownie przy każdym mnożeniu w obrębie kafelka.

Dla kafelka wyniku T × T jądro ładuje około 2·T·k elementów i wykonuje 2·T²·k FLOP-ów, więc intensywność rośnie z T. Większe kafelki to więcej ponownego użycia, aż do granicy pojemności pamięci na chipie. Sztuka pisania szybkiego jądra GEMM to w większości dobór kształtów kafelków pod konkretny chip i nakładanie ładowania następnego kafelka na obliczenia dla bieżącego.

Rdzenie tensorowe: obliczenia macierzowe w jednej instrukcji

Zwykły rdzeń GPU wykonuje jedno połączone mnożenie z dodawaniem (fused multiply-add) na tor (lane) w każdym cyklu. Rdzenie tensorowe NVIDIA, wprowadzone wraz z architekturą Volta w 2017 roku, wykonują zamiast tego w jednej instrukcji mały iloczyn macierzy. W architekturze Volta każdy rdzeń tensorowy liczył w każdym takcie zegara D = A × B + C na kafelkach 4 × 4, z wejściami w FP16 i akumulacją w FP32. Kolejne generacje poszerzały kafelki i dodawały formaty: Ampere wprowadził TF32 i BF16, Hopper FP8, a Blackwell FP4. Hopper dodał też instrukcje na poziomie grupy warpów (warpgroup) i sprzętową jednostkę (TMA) do asynchronicznego przenoszenia kafelków do pamięci współdzielonej.

Nie chodzi o konkretne kształty. Przeznaczenie krzemu na gęstą pętlę wewnętrzną matmul daje dużo więcej FLOP-ów na wat i na milimetr kwadratowy niż tory ogólnego przeznaczenia. W każdym współczesnym GPU do AI większość reklamowanych FLOP-ów to FLOP-y rdzeni tensorowych, a jądro, które z nich nie korzysta, zostawia większość chipa bezczynną.

Macierze systoliczne: podejście TPU

Pierwszy układ TPU firmy Google (Jouppi i in., 2017) posunął ten pomysł dalej dzięki macierzy systolicznej: siatce 256 × 256 złożonej z 65 536 8-bitowych jednostek mnożenia z akumulacją (multiply-accumulate). Wagi pozostają nieruchomo w komórkach, aktywacje wpływają z jednej strony, a sumy częściowe płyną przez siatkę, przesuwając się o jedną komórkę na takt, „jak krew pompowana przez serce”, i stąd pochodzi nazwa. Dane są używane ponownie, gdy przechodzą między sąsiednimi komórkami, zamiast wracać do pamięci. Późniejsze TPU i wiele układów do inferencji stosują tę samą zasadę w różnej skali.

Precyzja ciągle maleje

Każde zmniejszenie rozmiaru elementu o połowę to o połowę mniej przesyłanych bajtów, dwa razy więcej elementów w linii cache i mniej więcej dwukrotnie większa przepustowość rdzeni tensorowych na sprzęcie, który dany format obsługuje. Historia formatów liczbowych w AI podąża tym śladem:

• FP32 → FP16: połowa bajtów, ale przy wąskim zakresie wykładnika FP16 w treningu łatwo o przepełnienie.
• BF16: 8-bitowy wykładnik z FP32 i krótka mantysa. Ten sam zakres co FP32 przy połowie rozmiaru, dziś domyślny format do treningu.
• FP8 (E4M3 / E5M2): używany do treningu na dużą skalę. DeepSeek-V3 był trenowany głównie w FP8 (opisuje to notatka o architekturze DeepSeek).
• 4-bitowe wagi do inferencji: zwykle zapisane w formacie blokowym i dekwantyzowane do 16 bitów w rejestrach tuż przed instrukcją rdzenia tensorowego. To dokładnie wzorzec połączonej dekwantyzacji z notatek o kwantyzacji i atencji.

Zasady, których się trzymam

1. Najpierw policz intensywność arytmetyczną. Powie Ci, czy optymalizować obliczenia, czy pamięć, zanim napiszesz choćby linijkę kodu jądra.
2. Znaj punkt załamania swojego GPU. Poniżej niego jedyna dźwignia, która działa, to mniej bajtów na FLOP (kwantyzacja, batching, fuzja).
3. Korzystaj z rdzeni tensorowych. Jeśli profil pokazuje operacje matmul na zwykłych potokach FP32, najpierw popraw typ danych albo kształty.

Źródła: Williams, Waterman i Patterson, Roofline (CACM 2009). Whitepapery NVIDIA o architekturach Volta, Ampere i Hopper. Jouppi i in., In-Datacenter Performance Analysis of a Tensor Processing Unit (ISCA 2017). Wartości szczytowe to specyfikacje producentów; intensywności i czasy są policzone ze wzorów powyżej.