CuTeDSL w Perplexity

W Perplexity korzystamy z autorskiego silnika Runtime-Optimized Serving Engine (ROSE) do obsługi modeli, począwszy od embeddingów, aż po modele LLM o bilionach parametrów. Nasze środowisko wykonawcze, stanowiące zaplecze wszystkich interfejsów API, w tym Sonar, Search oraz Embeddings, zostało stworzone, aby umożliwić badaczom odkrywanie najodpowiedniejszych modeli i architektur, przy jednoczesnym szybkim wdrażaniu najlepszych kandydatów do środowiska produkcyjnego. W tym celu musimy dysponować możliwością wdrażania, eksploatacji oraz optymalizacji wnioskowania dla najnowocześniejszych architektur w szybkim tempie. Na łamach tego bloga przybliżamy, w jaki sposób jądra CuTeDSL zintegrowane z naszym silnikiem wnioskowania pozwalają nam budować wymagane jądra GPU, aby zapewnić modelom najwyższą wydajność na układach NVIDIA.

Niestandardowe modele, tworzone przez różne zespoły w Perplexity, są hostowane wewnętrznie na układach GPU NVIDIA Hopper oraz Blackwell przy użyciu silnika wnioskowania ROSE, zbudowanego przez zespół AI Inference. Silnik ROSE powstał początkowo w celu lepszej obsługi dostosowanych modeli Llama, zarówno na potrzeby dekodowania języka, jak i klasyfikacji, z interfejsem kompatybilnym z serwerem NVIDIA Triton Inference Server. Z czasem przekształcił się w pełnoprawny silnik wnioskowania, obsługujący znaczne modele LLM, a także ranking, klasyfikację, punktację oraz inne modele oparte na transformerach, dla szerokiego grona zespołów wspierających produkty Perplexity. Celem projektowym ROSE jest prostota: z biegiem czasu inwestowane są zasoby w osiągnięcie szczytowej wydajności, jednak głównym priorytetem pozostaje adaptacja do stale ewoluujących architektur oraz przypadków użycia naszych modeli.

Przestrzeń konfiguracji jądra obejmująca typy danych, wymiary głowy, stałe czasu kompilacji i ustawienia uruchamiania.

W ogólnym zarysie ROSE składa się z silnika, który dostosowuje model do interfejsu, z którym mogą współpracować klienci. Silniki są wysoce reużywalne, obsługują harmonogramowanie i przetwarzanie wsadowe żądań oraz pobierają dekodowane tokeny, embeddingi lub wyniki z modeli bazowych. Inicjalizacja urządzeń, kanały komunikacji międzyprocesowej i międzywęzłowej oraz ładowanie wag są w pełni obsługiwane przez nie. W przypadku modeli LLM silnik odpowiada za przetwarzanie wsadowe, fragmentowanie oraz próbkowanie. Odpowiada również za alokację pamięci KV, zarówno dla pełnej, jak i liniowej uwagi, a także dopasowywanie prefiksów. Silniki embeddingów są prostsze, gdyż wykonują jedynie przetwarzanie wsadowe on-line. Wykonanie jest zoptymalizowane w taki sposób, że silnik przygotowuje kolejną partię prac podczas równoległego działania modelu na akceleratorze.

ROSE udostępnia bogaty zbiór niestandardowych warstw do modelowania, opakowujących jądra GPU zbudowane za pomocą CuTeDSL, Triton, CUDA, CUTLASS oraz cuBLAS. W przypadku implementacji mnożenia macierzy i uwagi polegamy głównie na jądrach NVIDIA z bibliotek CUTLASS i cuBLAS. Pozostała część potoku, w tym embedding, norm, MoE oraz jądra aktywacji, stanowi część biblioteki jąder ROSE. Aby jak najlepiej wspierać sprzęt, na którym uruchamiamy wnioskowanie, ROSE wyposażono we własne prymitywy komunikacji międzyurządzeniowej dla routingu, dysponowania i łączenia MoE. Chociaż pozostałe jądra są stosunkowo proste, muszą być dokładnie wyspecjalizowane, aby osiągnąć szczytową wydajność na wszystkich różnych ukrytych wymiarach i urządzeniach, na których pracują. Ponadto polegamy na specjalizacji w czasie kompilacji, aby uwzględnić wszystkie parametry i modyfikacje różnych modeli (wagi z biasem lub bez biasu dla normy RMS, funkcje aktywacji itp.). W rezultacie zdecydowaliśmy się użyć języków DSL do ich implementacji: początkowo zbudowaliśmy ROSE przy użyciu Triton. Jednak od tego czasu przyjęliśmy CuTeDSL jako główne środowisko do programowania GPU.

Z matematycznego punktu widzenia jądra wnioskowania są zazwyczaj nieskomplikowane. W rezultacie, podczas tworzenia jądra, złożoność zazwyczaj nie leży w wyrażeniu tego, jakie obliczenia mają zostać wykonane, ale w tym, jak dokładnie sprzęt powinien to zrobić. Aby zbudować najlepsze jądra, potrzebujemy języka, który oferuje właściwe abstrakcje nad podstawowym sprzętem. Odpowiedni język umożliwia nam określenie operacji wysokiego poziomu w celu wykonania obliczeń, agresywnie je specjalizując i optymalizując we wszystkich punktach szerokiej przestrzeni konfiguracyjnej.

CuTeDSL to język DSL oparty na Pythonie dla procesorów graficznych NVIDIA, zbudowany w oparciu o algebrę układów CuTe oraz MLIR. Kompiluje się w trybie just-in-time do wysoce zoptymalizowanego kodu PTX, zachowując jednocześnie kontrolę nad niskopoziomowymi prymitywami sprzętowymi. Ma na celu oferowanie tych samych funkcji i wydajności co CUTLASS C++, bez kosztownych czasów kompilacji, które utrudniają szybkość programowania. Spośród języków, które przeanalizowaliśmy lub z którymi pracowaliśmy (w tym CUDA, cuTile, TVM IR oraz Triton), uważamy go za najlepiej dopasowany do tworzenia naszych jąder.

Podsumowując, choć nie jest tak wysokopoziomowy i zwięzły jak Triton lub cuTile, uznaliśmy, że algebra układów CuTe i jej kompozycje (udostępnione poprzez cute.Tensor) są rozsądnymi substytutami. Ponadto CuTeDSL oferuje nieograniczony dostęp do prymitywów sprzętowych, zapewniając nam poziom kontroli wymagany do osiągnięcia szczytowej wydajności we wszystkich przypadkach brzegowych. W przypadku CUDA możliwości specjalizacji w czasie kompilacji były ograniczone, ponieważ każdy dodatkowy parametr dodawał nieuzasadniony koszt do czasów kompilacji lub złożoność systemu budowania. W przeciwieństwie do tego, dzięki skompilowanemu w czasie rzeczywistym językowi DSL, możemy specjalizować się bardziej agresywnie bez poświęcania tempa iteracji. W odróżnieniu od TVM IR, który wykorzystywał CUDA jako krok pośredni, uznaliśmy, że potok oparty na MLIR, który redukuje do PTX, jest bardziej solidny i łatwiejszy do debugowania niż ten, który generował CUDA jako krok pośredni. Dodatkowo, dzięki użyciu niestandardowych operacji do opakowania instrukcji PTX, odkryliśmy, że możemy w pełni wykorzystać wszystkie funkcje bazowego sprzętu.

Większość kosztów kompilacji jąder CUDA przypada zazwyczaj na fazę rozwijania szablonów jąder, ponieważ optymalizacja i kompilacja od etapu PTX mogą zostać ukończone w ciągu kilku sekund. Przy wielu parametrach czasu kompilacji etap ten może wydłużyć się do kilku minut w przypadku niebanalnego jądra. To wąskie gardło jest potęgowane przez fakt, że możliwości kompilacji równoległej pojawiają się dopiero później, na poziomach IR i PTX, podczas gdy obsługa szablonów C++ odbywa się sekwencyjnie. Choć istnieją triki pozwalające na rozwiązanie tego problemu, poprzez przeniesienie instancji szablonów do oddzielnych jednostek kompilacji, wymagają one znaczących poprawek w systemach budowania i użycia nieporęcznych makr. Z kolei jądra CuTeDSL skompilowane w trybie JIT (Just-In-Time) nie wymagają skomplikowanych reguł budowania i nie utrudniają czasów kompilacji, ponieważ cykl programistyczny może skupiać się na jednym punkcie w przestrzeni konfiguracyjnej, podczas gdy staranne testy jednostkowe mogą później pokryć całą przestrzeń projektową jądra.

Mimo że kompilacja JIT przenosi koszty z czasu programowania na czas wykonania, narzut jest więcej niż akceptowalny. Gdy model jest wdrażany, aktywuje tylko kilka punktów z całej przestrzeni konfiguracyjnej, generując sekundowe narzuty podczas uruchamiania/rozgrzewania silnika wnioskowania. Dla odmiany wstępna kompilacja wszystkich konfiguracji jądra mogłaby trwać wiele minut. Ponadto narzuty można dodatkowo ograniczyć, buforując artefakty skompilowane w JIT.

Wysoki koszt kompilacji CUDA jest potęgowany przez trudności w śledzeniu źródła błędów, gdy stosowane jest wysyłanie i specjalizacja oparte na szablonach, jak ma to miejsce w CuTe i CUTLASS. Dzięki DSL błędy mogą być wskazywane w wierszu kodu, który je wygenerował, co znacząco ułatwia debugowanie. Ponadto znacznie prościej jest zmapować wygenerowany kod do instrukcji DSL, z której pochodzi, poprzez pośrednie etapy MLIR lub polegając na informacjach debugowania dołączonych do wygenerowanego PTX. Dzięki tym aspektom DSL odkryliśmy, że jesteśmy znacznie bardziej produktywni podczas tworzenia jąder.

Choć jądra wykazują wiele podobieństw między modelami, muszą być wyposażone w szereg parametrów konfiguracyjnych, które mogą mieć wpływ na wydajność. Odpowiednie parametry obejmują:

Wymiar ukryty / Wymiar głowy / Liczba głów uwagi

Typy danych wejściowych/wyjściowych

Schemat kwantyzacji

Stałe czasu kompilacji (biasy, funkcje aktywacji itp.)

Architektura ROSE, w tym jej silnik, model, biblioteka jądra i usługi.

Specjalizacja w czasie kompilacji jest przydatna zarówno dla wydajności, jak i debugowania. Szczególnie w procesorze potokowym, jak ma to miejsce w przypadku GPU, nieistniejąca gałąź jest lepsza niż zawsze zajęta lub nigdy niezajęta gałąź. Ponadto kompilacja jądra dla określonego zestawu parametrów konfiguracyjnych skutkuje mniejszą liczbą instrukcji PTX/SASS, które są łatwiejsze do analizy i debugowania przy użyciu cuda-gdb podczas dochodzenia.

Dla danego rozmiaru wsadu i konfiguracji można oszacować oczekiwany ruch do pamięci DRAM oraz liczbę operacji arytmetycznych. Następnie można podejmować optymalne decyzje w każdym przypadku, aby uzyskać optymalną alokację siatki, synchronizację oraz wykorzystanie niektórych funkcji sprzętowych. CuTeDSL pozwala nam definiować i ponownie wykorzystywać komponenty do głębokiej specjalizacji bez ponoszenia niedopuszczalnego narzutu czasu kompilacji.

Ogólnie rzecz biorąc, wydajność jąder wnioskowania musi być zoptymalizowana dla dwóch oddzielnych przypadków użycia:

Prefill: mały rozmiar wsadu (kilka fragmentów, sekwencji), wiele tokenów na sekwencję. Przepustowość ma być zmaksymalizowana zarówno dla embeddingów, jak i tokenów wejściowych modeli LLM. Opóźnienie rzędu od dziesiątek do setek milisekund.

Decode: rozmiary wsadu od małych do dużych, od jednego (brak dekodowania spekulatywnego) do kilku tokenów na sekwencję (spekulacja draft-target, przewidywanie wielu tokenów). Opóźnienie bliskie opóźnieniu uruchomienia jądra.

Dla GEMM i uwagi etap prefill działa na dużej liczbie tokenów, a wydajność jest ograniczona przez przepustowość obliczeniową urządzenia. Dekodowanie wiąże się z tą samą operacją przy małych rozmiarach wsadu, gdzie czasy wykonania są zdominowane przez koszt ładowania aktywacji i wag z pamięci DRAM do strumieniowych procesorów wielordzeniowych i rdzeni tensorowych, ograniczony przez przepustowość pamięci urządzenia. Mimo że rozmiar wsadu lub całkowita liczba tokenów w wsadzie jest stosunkowo wyraźnym wskaźnikiem, próg przełączania między prefill i decode jest bardziej subtelny, ponieważ przepustowość pamięci jest również determinowana przez wymiar ukryty modelu i typ danych aktywacji. Jądra są zatem dzielone na implementacje prefill i decode, jeśli różnice algorytmiczne lub przypisania siatki/bloku dawałyby przewagę jednemu z nich.

Parametrem, który jest uniwersalnie różny między jądrami prefill a decode, jest liczba warpów w bloku (liczba wątków podzielona przez rozmiar warpa wynoszący 32 wątki na warp). Pożądane jest, aby ten parametr był stałą czasu kompilacji, ponieważ może pomóc w alokacji rejestrów. W CUDA odbywa się to zazwyczaj poprzez określenie liczby wątków za pomocą __launch_bounds__ i podanie liczby warpów jako parametru szablonu:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

Liczba rejestrów dostępnych dla wszystkich wątków w bloku jest ograniczeniem sprzętowym, wynoszącym 64K na układach Hopper i Blackwell. Jeśli liczba wątków jest znana z wyprzedzeniem, kompilator dokładnie wie, ile rejestrów może przydzielić każdemu wątkowi i może skutecznie kompilować bardziej agresywnie zoptymalizowane jądra. Jednak wiąże się to z kosztem czasu kompilacji, ponieważ teraz jądro musi zostać ponownie skompilowane dla różnych rozmiarów warpów. Dzięki CuTeDSL nie jest to już problemem, ponieważ nie musimy ponosić kosztu kompilacji dla całej przestrzeni konfiguracyjnej z góry. DSL oczekuje również, że rozmiar bloku będzie stałym wyrażeniem, domyślnie wykonując optymalizację sugerowaną przez __launch_bounds__ w CUDA.

Oprócz uczynienia liczby warpów stałą, pożądane jest różnicowanie jej między prefill a decode. Aby w pełni wykorzystać przepustowość pamięci GPU, nie wystarczy zawsze przypisywać dużej liczby warpów na blok, ponieważ pełną przepustowość można osiągnąć jedynie poprzez wykorzystanie jak największej liczby dostępnych strumieniowych procesorów wielordzeniowych (SM). W konsekwencji, dla danej liczby tokenów oraz znanych typów danych i wymiarów ukrytych, najpierw wybieramy idealny wymiar wektora na wątek (8 dla bfloat16, 4 dla float32), aby zapewnić, że wygenerowany kod używa 128-bitowych instrukcji ładowania/zapisu (256-bitowych w przypadku Blackwell). Znając rozmiar wektora, możemy obliczyć liczbę warpów wymaganą przez jądro. Najpierw staramy się zmaksymalizować liczbę jednostek SM: jeśli warpów jest mniej niż jednostek SM, wybierany jest pojedynczy warp na blok. Jest to zazwyczaj wariant dekodowania jądra. Jeśli warpów jest wystarczająco dużo, stopniowo zwiększamy ich liczbę do 8, 16 lub 32, wyznaczając górną granicę, po której zazwyczaj nie widzimy już poprawy wydajności.

Choć Triton wykonywał niektóre z tych optymalizacji wewnątrz, odkryliśmy, że możemy osiągnąć podobne wyniki poprzez staranne wykorzystanie algebry układów CuTe, nie będąc ograniczonymi przez niedociągnięcia kompilatora wysokiego poziomu. Zwłaszcza w przypadku redukcji horyzontalnych sub-warp, uznaliśmy, że CuTeDSL oferuje nam drobnoziarnistą kontrolę wymaganą do wygenerowania optymalnego kodu.

W niektórych sytuacjach, aby rozdzielić pracę między jak największą liczbę bloków, można użyć cooperative_groups::this_grid().sync() w celu zsynchronizowania wszystkich uczestniczących bloków. Obowiązują ograniczenia: ponieważ blok nie może zostać wyłączony z harmonogramu SM po rozpoczęciu wykonywania, aby zagwarantować postęp, liczba bloków nie może przekraczać liczby jednostek SM dostępnych na urządzeniu, aby zapewnić, że w dowolnym momencie wszystkie bloki mogą aktualizować i odpytywać lokalizacje pamięci globalnej zaangażowane w synchronizację. To ograniczenie jest egzekwowane przez cudaLaunchCooperativeKernel. Koszt bariery siatki wynosi około 3us na układzie Hopper, plus koszt opróżnienia wszystkich poprzednich operacji pamięci poprzez niejawny barierę pamięci w niektórych sytuacjach. Ogólnie rzecz biorąc, jest to porównywalne z narzutem podzielenia jądra na dwa uruchomienia, co prowadzi do decyzji projektowej wrażliwej na wydajność. Pomimo faktu, że Programmatic Dependent Launch (PDL) może zmniejszyć narzuty uruchomienia od układów Hopper wzwyż, dla jąder dekodowania unikanie narzutu jest opłacalne.

Jeśli jedno z obciążeń wymaga więcej bloków niż liczba jednostek SM, jądro powinno zostać podzielone, polegając na niejawnej barierze uruchomień jądra. Jest to głównie przypadek dla prefill. Jeśli liczba bloków idealna dla wszystkich obciążeń mieści się w siatce, preferowana jest bariera siatki, jak ma to miejsce w przypadku jąder dekodowania działających na mniejszej liczbie tokenów.

W CuTeDSL decyzje te można podejmować w kodzie uruchomieniowym po stronie hosta, a jądra mogą być komponowane i przeplatane barierami siatki w miarę potrzeb, co zilustrowano na poniższym przykładzie. Dzięki takiej strategii hybrydowej jądra dekodowania można zoptymalizować w celu osiągnięcia minimalnych opóźnień, a jądra prefill w celu uzyskania maksymalnej przepustowości. W przeciwieństwie do rozwiązań opartych na CUDA, logikę tę można zdefiniować bardziej zwięźle i z mniejszym narzutem czasu kompilacji.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

W przypadku naszych jąder MoE all-to-all, jak wyszczególniono poniżej, elastyczność w budowaniu jąder z barierami siatki lub bez nich oraz optymalizacja ich w czasie kompilacji pomogły nam osiągnąć znaczne zyski wydajnościowe dzięki CuTeDSL.

Pierwszym jądrem, które zbudowaliśmy w CuTeDSL, było wyspecjalizowane jądro normy RMS, zoptymalizowane pod kątem małych wymiarów. Norma QK, stosowana w warstwie self-attention w najnowszych modelach Qwen, jest stosowana na poziomie poszczególnych głów, dla 64, 128 lub 256 elementów. W przeciwieństwie do zwykłej normy RMS, która operuje na ukrytych wymiarach przekraczających tysiące, wymagając horyzontalnej redukcji zarówno poprzez pamięć współdzieloną, jak i rejestry, norma QK wymaga redukcji jedynie pomiędzy warpami, z możliwością ściśnięcia wielu głów w pojedynczy warp. Stało się to przypadkiem użycia, w którym mogliśmy agresywnie optymalizować w czasie kompilacji, aby znaleźć najlepsze przypisania token-do-warpa dla wszystkich wymiarów głów i typów danych.

Norma RMS jest dość prostą operacją, choć wymaga sumy horyzontalnej do obliczenia normy kwadratowej wektora:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

Ponieważ głowy uwagi są zazwyczaj małe, a operacje na wątkach są silnie wektoryzowane, pojedynczy warp może obliczyć normę kwadratową. Aby zmaksymalizować przepustowość pamięci, każdy wątek powinien przetwarzać 128 bitów danych wejściowych, co przekłada się na 8 wartości bfloat16. W konsekwencji pół-warp jest wystarczający do przetworzenia głowy, a normę kwadratową można wyznaczyć za pomocą redukcji horyzontalnej wykonanej przez shuffle_sync_bfly. Zaczynając od wymiaru głowy i typu danych, możemy określić liczbę głów na warp oraz określić argumenty i liczbę rund redukcji w czasie kompilacji.

Następnie możemy zbudować wyspecjalizowane jądro za pomocą stosunkowo niewielkiej liczby linii kodu, polegając na automatycznie wektoryzowanych ładowaniach/zapisach i używając niskopoziomowych operacji specyficznych dla architektury do redukcji warp. Jądro określa, nad którą głową (Q lub K) ma pracować i identyfikuje odpowiednie wagi do zastosowania. Ładuje cały wektor z głowy i najpierw sumuje lokalnie. Następnie, w oparciu o liczbę wątków na głowę, obliczana jest norma kwadratowa. Na koniec norma RMS, automatycznie wektoryzowana, jest wyznaczana i zapisywana z powrotem do pamięci. Parametry specyficzne dla modelu, takie jak WEIGHT_BIAS oraz EPS, są obliczane jako stałe (constant-folded).

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

W porównaniu do zwykłego jądra normy RMS, ten wyspecjalizowany wariant zapewnia 2-3x lepszą przepustowość i opóźnienie zarówno dla prefill, jak i decode.

Dysponowanie (dispatch) i łączenie (combine) MoE przetasowują tokeny wejściowe do ekspertów, do których są kierowane, wykonując ważoną akumulację po tym, jak jądra MoE GEMM wykonają odpowiednie mnożenie macierzy. Wydajność jest szczególnie interesująca, ponieważ różne produkty polegają na modelach MoE różnych rozmiarów. W oparciu o typy wdrożenia wyróżniamy różne jądra do optymalizacji:

TP=1: Lokalne jądro dla pojedynczego GPU. Dysponowanie przetasowuje tokeny do odpowiednich ekspertów, a łączenie je kumuluje. Używane głównie dla wdrożeń Qwen3 30B na H100.

DP=1, TP ≠ 1: Równoległość tensorowa, jądro jednowęzłowe. Dane wejściowe są replikowane, dysponowanie wybiera tokeny dla bieżącej rangi, a łączenie kumuluje je pomiędzy wszystkimi rangami. Akumulacja poprzez NVLink lub NVLS. Używane dla dużych prefillerów.

DP ≠ 1, NVLink: Jądro zrównoleglenia danych z opcjonalną równoległością tensorową. Każda ranga wysyła tokeny do każdej innej rangi hostującej określonych ekspertów. Informacje o routingu są agregowane, tokeny są przesyłane poprzez NVIDIA NVLink. Używane dla wdrożeń H200 lub GB200 zrównoleglonych danymi.

DP ≠ 1, InfiniBand: Jądro zrównoleglenia danych i równoległości tensorowej Peer-to-Peer działające zarówno przez NVLink, jak i NVIDIA Quantum InfiniBand. Używane dla wielowęzłowych wdrożeń Hopper.

Lokalne jądro zlicza liczbę tokenów przypisanych do każdego eksperta w oparciu o tablicę routingu i przypisuje unikalny indeks każdemu tokenowi w ramach listy tokenów należących do każdego eksperta. Po kroku routingu tokeny są kopiowane do opcjonalnie dopełnionych buforów, aby zostać przekazane do jądra MoE GEMM. Używając lokalizacji tokenów, jądro łączenia wykonuje następnie akumulację w rejestrach, zapisując końcowe tokeny do bufora wyjściowego.

Faza dysponowania różni się między prefill a decode: dla kilku tokenów istnieje niewiele informacji o routingu do odczytania, więc przetwarzamy je w pojedynczym bloku i sygnalizujemy resztę siatki, przy czym poszczególne bloki kopiują fragmenty tokenów. Przypisanie tokenów jest określane poprzez atomowe zwiększanie liczników w pamięci współdzielonej. Ponieważ prefill musi przetworzyć więcej informacji o routingu, routing jest dzielony między wiele bloków w osobnym jądrze. Indeksy są najpierw przypisywane lokalnie przy użyciu liczników w pamięci współdzielonej, a następnie agregowane między wszystkimi blokami w pamięci globalnej, aby określić pozycję każdego tokenu. Siatka tego jądra jest ograniczona do liczby jednostek SM, aby móc używać barier siatkowych. Kolejne jądro wykonuje kopie, uruchamiane na znacznie większej liczbie bloków i warpów, aby w pełni wykorzystać dostępną przepustowość pamięci.

Jądro łączenia jest stosunkowo proste, ponieważ wykonuje wektoryzowaną akumulację, polegając na informacjach o routingu i przesunięciach obliczonych przez dysponowanie. CuTeDSL okazał się tutaj wartościowy, ponieważ obliczenia można było wyrazić poprzez proste tensory CuTe, które kompilują się do wydajnego kodu dla wszystkich typów danych i wymiarów ukrytych.

W wdrożeniu z równoległością tensorową informacje o routingu i tokeny wejściowe są replikowane. Lokalne jądro dysponowania może być w pełni ponownie wykorzystane z niewielką zmianą, aby kierować i wybierać tylko tokeny przypisane do ekspertów znajdujących się na bieżącej randze. Z pewnymi parametrami i specjalizacją czasu kompilacji lokalne jądra dysponowania są ponownie wykorzystywane.

Ponieważ liczba rang jest zazwyczaj mniejsza niż liczba ekspertów, do których kierowany jest token, w przeciętnym przypadku można założyć, że każdy token ma co najmniej jedną ze swoich replik skierowaną do każdej rangi. W tej sytuacji lokalne jądro łączenia może być ponownie wykorzystane do agregacji lokalnych kopii tokenu, po czym następuje all-reduce w celu zsumowania między uczestniczącymi rangami. W przypadku jądra prefill uruchamiamy oddzielne jądra, jednak decode łączy obie operacje w pojedyncze uruchomienie.

Dla all-reduce używamy zwykłego NVLink, jeśli uczestniczą tylko dwie rangi, przełączając się na NVLS, gdy zaangażowane są 4 lub więcej rang, co zapewnia wzrost przepustowości o 10-20%. Dzięki DSL możemy współdzielić większość logiki redukcji i wymiany buforów peer między MoE a zwykłymi jądrami all-reduce. Jest to sytuacja, w której ścisła kontrola nad sprzętem jest korzystna, ponieważ operacje DSL nad multimem.ld_reduce oraz multimem.st są wymagane dla NVLS. Implementacja samej redukcji w dużej mierze podąża za przykładem all_reduce z testów NCCL.

Jądra zrównoleglenia danych (Data Parallel) polegają na NVLink w celu agregacji tokenów dla wielu urządzeń. W fazie dysponowania, przed wymianą tokenów, rangi najpierw kierują lokalnie, aby zidentyfikować docelowe przesunięcie na każdej randze dla każdego wysyłanego tokenu. Następnie zaczynają kopiować tokeny do buforów odbiorczych przydzielonych dla każdego nadawcy, aby uniknąć konfliktów. Zapisy typu release-store na token są używane do sygnalizowania zakończenia. Po stronie odbiorcy informacje o routingu są agregowane, aby móc skondensować wszystkie tokeny w jednym buforze dla MoE GEMM. Odbiorcy odpytują flagę zapisaną przez rówieśnika za pomocą acquire-load, aby poczekać na zakończenie zapisów.

Po stronie łączenia tokeny są kopiowane przez NVLink do rang źródłowych, z flagą na token wskazującą odbiór. Jądro łączenia jest podobne do wariantów lokalnego i jednowęzłowego, z niewielkimi różnicami w adresowaniu i pewną dodatkową logiką synchronizacji.

Zbudowaliśmy bibliotekę opakowań PTX, aby udostępnić ładowania i zapisy typu release-acquire o zasięgu systemowym i urządzeniowym w celu zaimplementowania tego jądra. Poprawiając ziarnistość flag synchronizacji i przydzielając więcej bloków do wysyłania i odbierania tokenów przez NVLink, osiągnęliśmy 2-krotną poprawę w porównaniu z naszymi poprzednimi jądrami CUDA.

Nasze jądra Peer-to-Peer osiągały wcześniej szczytową wydajność na układach GPU Hopper w połączeniu z adapterami InfiniBand ConnectX-7. Od tego czasu przenieśliśmy jądra łączenia do CuTeDSL w celu uzyskania dalszych zysków wydajności. Ponieważ wersje CUDA już w dużej mierze polegały na specjalizacji w czasie kompilacji, okazały się one wąskimi gardłami czasu kompilacji w naszym repozytorium, zmuszając nas do pominięcia niektórych parametrów. Migrując do kompilowanego w JIT języka DSL, wyspecjalizowaliśmy się pod kątem większej liczby parametrów i dostosowaliśmy alokacje siatki, wprowadzając dalsze różnice między prefill a decode dla około 3% wzrostu opóźnienia.

Z perspektywy jąder ROSE wsparcie dla Blackwell dzieli się na dwie kategorie: jądra agnostyczne względem architektury oraz jądra specyficzne dla architektury. Jądra specyficzne dla architektury, polegające na rdzeniach Tensor, wymagają wyraźnej pracy adaptacyjnej. Jednak dzięki użyciu DSL, jądra agnostyczne względem architektury adaptują się bezproblemowo.

GPU Blackwell posiada do 160 jednostek SM, w przeciwieństwie do 132 w układzie Hopper. Dla większości jąder jest to stała czasu kompilacji, a logika w kodzie uruchomieniowym po stronie hosta uwzględnia to przy określaniu optymalnego jądra (prefill vs decode) do uruchomienia. Ponadto Blackwell wprowadza 256-bitowe instrukcje ładowania, które mogą być łatwo emitowane poprzez zwiększenie szerokości wektora (ostatni wymiar tensorów CuTe) w celu poprawy wykorzystania przepustowości pamięci.

Blackwell wprowadza sprzętowe wsparcie dla znaczącej liczby schematów kwantyzacji (MXFP8, MXFP4, NVFP4). Wsparcie dla nich zostało zaimplementowane głównie poprzez modyfikację istniejących jąder dla skalowania blokowego FP8. Wsparcie JIT jest tutaj kluczowe, ponieważ nowe tryby mogą być obsługiwane bez eksplozji przestrzeni konfiguracyjnej czasu kompilacji. W przeszłości uważaliśmy Triton za niewygodny w wyrażaniu redukcji horyzontalnych dla kilku elementów, powyżej połowy lub ćwiartki warpa. Bardziej drobnoziarnista kontrola nad redukcjami była kluczowa dla osiągnięcia optymalnej wydajności podczas obliczania wartości maksymalnej bezwzględnej dla 16 lub 32 elementów odpowiednio dla NVFP4 i MXFP8/MXFP4.

Jako stosunkowo nowy język, LLM miały niewielką styczność z kodem CuTeDSL podczas swojego uczenia wstępnego. Ponadto styl naszego obecnego kodu, z jądrami zoptymalizowanymi pod kątem aktywacji i norm, był nieco inny od istniejących przykładów CUTLASS lub Quack, które były silnie zoptymalizowane pod kątem TMA i specyficznych dla architektury instrukcji mnożenia macierzy. Osiągnięcie punktu, w którym LLM mogłyby pomagać w pisaniu jąder, było procesem stopniowym.

Ręcznie zbudowaliśmy początkowy korpus jąder i funkcji użytkowych, przestrzegając spójnych wzorców, zanim mogliśmy polegać na agentach w kwestii konwersji jąder. W miarę rozwoju korpusu mogliśmy przeprowadzać zgrubne tłumaczenia z CUDA na CuTeDSL, po których następowała znacząca ręczna poprawa. W szczególności odkryliśmy, że modele zmagały się z konwersjami wskaźnik-na-tensor oraz wektoryzacją, utykając na próbach naprawy błędów typów, gdy celem było skompilowanie testu. Wykazały znikome zrozumienie opartej na siatce natury jąder, myląc synchronizację wątków i siatki. Co najbardziej frustrujące, istniała pokusa halucynowania nieistniejących tłumaczeń bibliotek specyficznych dla CUDA zamiast budowania w oparciu o nasze własne zamienniki CuTeDSL. Z biegiem czasu, gdy korpus jąder CuTeDSL rósł, odkryliśmy, że LLM potrafią zbudować większość kodu szkieletowego i wypełnić luki, pozwalając nam skupić się na projekcie wysokiego poziomu.

W porównaniu do CUDA uważamy, że CuTeDSL zapewnia znacznie lepszą wydajność w czasie kompilacji, co pozwala nam specjalizować się pod kątem większej liczby parametrów i osiągać poprawę wydajności. Choć jądra są nieco bardziej rozwlekłe ze względu na brak niektórych abstrakcji w porównaniu do Triton, uważamy, że złożoność jest uzasadniona, ponieważ drobnoziarnista kontrola nad sprzętem pozwala nam budować szybsze jądra.

W momencie pisania tego tekstu różnica w wydajności między jądrami CUTLASS C++ a CuTeDSL jest minimalna, co otwiera ekscytujące możliwości komponowania i łączenia wysokowydajnych jąder GEMM z komunikacją międzyurządzeniową, normalizacją czy aktywacjami. Dzięki zaangażowaniu w ten język DSL będziemy w stanie wydobyć z naszych urządzeń dalszą wydajność, aby poprawić przepustowość naszych obciążeń, od embeddingów po modele LLM.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Wnioskowanie w Perplexity AI

Dlaczego CuTeDSL?

Doświadczenie programisty

Duża przestrzeń konfiguracyjna

Wstępne wypełnianie i dekodowanie specjalizacji

Synchronizacja siatki

Jądra CuTeDSL

Norma QK

Dysponowanie/Łączenie Mixture-of-Experts

Lokalne

Jednowęzłowe

Zrównoleglenie danych

Peer-to-Peer

Wsparcie dla Blackwell

Wspomagane przez AI generowanie jąder

Wnioski i przyszłe prace

Bibliografia

Cytat