Szybkie jądra uwagi i fuzja operatorów reorganizują obliczenia w transformerach, aby ograniczyć przemieszczanie danych i uruchamianie jąder na akceleratorze. Te optymalizacje w zasadzie pozostawiają niezmienione wagi modelu i jego funkcję matematyczną, choć kolejność operacji i precyzja mogą powodować niewielkie różnice numeryczne. Opublikowany przez Machine Learning Mastery tutorial szczegółowo opisuje te techniki, wraz z przykładami kodu.
Warstwa transformer opiera się na typowych operacjach: projekcjach liniowych, uwadze, normalizacji, funkcjach aktywacji i dodawaniach resztkowych. Bezpośrednia implementacja uruchamia osobne jądro obliczeniowe, zwane „kernel », dla każdej operacji i zapisuje każdy tensor pośredni w pamięci akceleratora. Obliczenia są poprawne, ale czas wykonania może być bardziej poświęcony na przemieszczanie wartości i uruchamianie jąder niż na użyteczne obliczenia.
Fuzja operatorów
Blok SwiGLU, używany przez wiele modeli z samym dekoderem, ilustruje ten problem. Bezpośrednia implementacja uruchamia mnożenia macierzy dla projekcji, jądro dla funkcji SiLU, kolejne dla mnożenia element po elemencie, a następnie ostatnie mnożenie macierzy. Na każdej granicy tensor pośredni może być zapisany w pamięci o wysokiej przepustowości (HBM), a następnie odczytany ponownie. Fuzja polega na połączeniu tych etapów w jedno jądro. Oszczędza czas, gdy przemieszczanie pamięci i uruchamianie jąder mają duży wpływ, a znacznie mniej, gdy duże mnożenie macierzy już dominuje w wykonaniu. Wybór jądra pozostaje empiryczny: trzeba profilować całe obciążenie.
Inne fuzje są powszechne w transformerach, jak dodanie połączenia resztkowego z następującą po nim normalizacją RMSNorm lub LayerNorm, lub połączenie dekwantyzacji z mnożeniem macierzy. Fuzja ma jednak swoje ograniczenia: operacje o niekompatybilnych organizacjach danych lub bardzo duży stan pośredni źle się łączą. Zbyt duże jądro może zwiększyć nacisk na rejestry, zmniejszyć stopień zajętości i stać się wolniejsze niż kilka dobrze dostrojonych jąder.
Uwaga typu FlashAttention i uwaga stronicowana
Klasyczna uwaga oblicza macierz wyników o rozmiarze T² na głowę, gdzie T to długość sekwencji. Na GPU jej wygenerowanie, zapisanie w pamięci globalnej i ponowne odczytanie zużywa dużo przepustowości, co może stać się wąskim gardłem podczas długiego wstępnego wypełniania („prefill »). FlashAttention przetwarza zapytania, klucze i wartości blokami. Dla każdego bloku oblicza część wyników, aktualizuje statystyki softmax i akumuluje odpowiadające wyjście w pamięci na układzie, bez przechowywania pełnej macierzy. PyTorch udostępnia tę optymalizację za pomocą funkcji scaled_dot_product_attention(). Tutorial wskazuje, że narzucenie backendu na sztywno może spowodować błąd kodu, jeśli żaden backend nie obsługuje danych wejściowych lub sprzętu.
Pamięć podręczna KV, która przechowuje już obliczone klucze i wartości, rośnie o jeden token na raz, do nieprzewidywalnej końcowej długości. Ciągła alokacja na zapytanie pozostawia więc zarezerwowaną i niewykorzystaną przestrzeń. Stronicowana pamięć podręczna dzieli przechowywanie na bloki o stałym rozmiarze i opiera się na tabeli bloków, która łączy pozycje logiczne z blokami fizycznymi. Jądro uwagi musi wtedy tłumaczyć każdą pozycję przez tę tabelę. Stronicowanie nie zmniejsza liczby elementów logicznych, ale ogranicza fragmentację i zarezerwowaną pojemność.
W podanym przykładzie vLLM używa modelu Qwen/Qwen2.5-1.5B-Instruct z blokami po 16 tokenów (block_size=16) i backendem „TRITON_ATTN ». Według tutoriala, vLLM sprawdza zgodność backendu z modelem i sprzętem podczas inicjalizacji i zgłasza błąd w przypadku niezgodności.
Fuzja i CUDA Graphs z torch.compile
Funkcja torch.compile() przechwytuje sekwencję operacji za pomocą TorchDynamo, a następnie kompiluje ją za pomocą TorchInductor, swojego domyślnego backendu. TorchInductor może wygenerować jedno jądro Triton dla kompatybilnych operacji, które tryb standardowy uruchomiłby osobno. Przerwania grafu, spowodowane na przykład przez przepływ sterowania zależny od danych lub konwersje tensorów na skalary Pythona, ograniczają obszary, które kompilator może zoptymalizować.
Tryb „reduce-overhead » nakazuje TorchInductor używać CUDA Graphs, gdy przechwycony graf się do tego kwalifikuje. CUDA Graph rejestruje sekwencję operacji GPU i odtwarza ją z pojedynczym zgłoszeniem do procesora. Nie tworzy dodatkowej fuzji, uzupełnia tę z TorchInductor. Wymaga stabilnych adresów pamięci i kompatybilnych kształtów oraz może zużywać więcej pamięci. Tutorial zaleca sprawdzanie wyników z tolerancjami odpowiednimi dla typu danych, profilowanie wykonania i mierzenie wydajności end-to-end.
