Przełom w optymalizacji GPU: AI tworzy kod CUDA wydajniejszy niż PyTorch

Sztuczna inteligencja wkracza na kolejny poziom, tym razem optymalizując kod niskopoziomowy. Zespół naukowców z Uniwersytetu Stanforda opracował metodę, w której modele językowe generują kernele CUDA – specjalistyczne programy działające bezpośrednio na procesorach graficznych Nvidia – osiągając rezultaty lepsze niż w przypadku ręcznie optymalizowanych rozwiązań z biblioteki PyTorch.

Kernele CUDA to fundament obliczeń w wielu aplikacjach AI, odpowiedzialne za operacje takie jak mnożenie macierzy czy przetwarzanie obrazów. PyTorch, framework stworzony m.in. przez Meta, oferuje bogaty zestaw gotowych funkcji GPU, niezbędnych w zadaniach uczenia maszynowego. Okazuje się jednak, że wygenerowane przez AI alternatywy mogą być jeszcze szybsze.

Przewaga AI w benchmarkach

W przeprowadzonych testach, kernele CUDA stworzone przez modele językowe znacząco przyspieszyły działanie w porównaniu do wbudowanych funkcji PyTorch. Przykładowo, jeden z nich zoptymalizował normalizację warstwy – proces standaryzacji wartości w sieciach neuronowych – aż 4,8-krotnie.

Poprawa wydajności była widoczna również w innych operacjach, takich jak konwolucja obrazu (Conv2D), funkcja softmax (przekształcająca wyniki w prawdopodobieństwa) oraz złożone operacje łączące konwolucję, aktywację ReLU i max-pooling. W wielu przypadkach, automatycznie wygenerowany kod okazał się szybszy od tego z PyTorch.

Równoległe poszukiwanie optymalizacji

Kluczem do sukcesu okazało się zastosowanie benchmarku KernelBench, w którym model językowy próbuje zastąpić konkretne operatory PyTorch własnymi kernelami CUDA, dążąc do maksymalnej wydajności GPU. Zespół wykorzystał modele OpenAI o3 oraz Gemini 2.5 Pro, które stosują strategie optymalizacji równoległej w wielu iteracjach. Każdy wygenerowany kernel był następnie sprawdzany pod kątem poprawności i szybkości działania.

Innowacyjne podejście badaczy ze Stanford polegało na wyrażaniu pomysłów optymalizacyjnych w języku naturalnym i generowaniu wielu wariantów kodu jednocześnie. Te warianty były wykonywane równolegle, a tylko najszybsze wersje przechodziły do kolejnej rundy, co umożliwiło eksplorację szerszego zakresu rozwiązań. Najefektywniejsze kernele wykorzystywały sprawdzone techniki, takie jak efektywniejszy dostęp do pamięci, nakładanie się operacji arytmetycznych i pamięciowych, redukcja precyzji danych (np. z FP32 do FP16), lepsze wykorzystanie jednostek obliczeniowych GPU oraz uproszczenie struktur pętli.

Implikacje dla przyszłości

Naukowcy podkreślają, że ich metoda nie zastępuje istniejących metod trenowania, lecz je uzupełnia. Proces poszukiwania optymalizacji generuje syntetyczne dane, które mogą posłużyć do trenowania przyszłych modeli. Wynikiem jest więc nie tylko szybszy kod w czasie wykonywania, ale także cenne dane treningowe do ulepszania modeli językowych.

Jednym z przykładów jest poprawa wydajności wygenerowanego kernela dla konwolucji obrazu (Conv2D), który po 13 iteracjach osiągnął niemal 180% szybkości PyTorch. Kluczowe ulepszenia obejmowały konwersję konwolucji na mnożenie macierzy wykorzystujące specjalistyczne rdzenie tensorowe GPU, użycie podwójnego buforowania w celu nałożenia obliczeń i dostępu do pamięci oraz wstępne obliczanie indeksów pamięci w pamięci współdzielonej dla szybszego pobierania danych.

Końcowa wersja zawierała zaawansowane techniki programowania CUDA, które zazwyczaj można zobaczyć tylko u programistów z wieloletnim doświadczeniem w optymalizacji GPU.

Wyzwana i ograniczenia

Zespół nadal pracuje nad pewnymi niedociągnięciami. Generowane przez AI kernele mają trudności z nowszymi zadaniami AI, które wykorzystują typy danych o niższej precyzji, takie jak FP16. W jednym teście kernel mnożenia macierzy FP16 osiągnął jedynie 52% szybkości PyTorch. Jeszcze gorzej wypadło Flash Attention, technika intensywnie wykorzystująca pamięć w dużych modelach językowych, gdzie kernel wygenerowany przez AI osiągnął zaledwie 9% tempa PyTorch.

Jednak naukowcy nie są zniechęceni. Zwracają uwagę, że wcześniej nikt nie potrafił wiarygodnie generować tego rodzaju kerneli automatycznie. Nawet przy ograniczonych zasobach na proces poszukiwania, nowe podejście już robi postępy.

Wyniki te są zgodne z innymi niedawnymi pracami pokazującymi, jak strategie poszukiwania równoległego, w połączeniu z potężnymi modelami językowymi, mogą tworzyć imponujące komponenty systemowe. Podobne rezultaty zaobserwowano w projektach takich jak AlphaEvolve DeepMind i funkcja Deep Think Gemini 2.5 Pro.