CuTeDSL u kompaniji Perplexity

U kompaniji Perplexity koristimo interni Runtime-Optimized Serving Engine (ROSE) za pokretanje modela, od onih za vektorska upoređivanja (embedding) do jezičkih modela sa bilionima parametara (LLM). Smešten iza svih API-ja, uključujući Sonar, Search i Embeddings, naš izvršni okruženje je napravljen da osnaži istraživače da istraže najpogodnije modele i arhitekture, a da pritom brzo primene najbolje kandidate u produkciju. Da bismo to postigli, moramo biti u mogućnosti da implementiramo, primenimo i optimizujemo inferenciju za najsavremenije arhitekture ubrzanim tempom. U ovom postu na blogu istražujemo kako su CuTeDSL kerneli integrisani u naš mehanizam za inferenciju omogućavaju nam da izgradimo potrebne GPU kernele kako bismo doveli modele do maksimalnih performansi na NVIDIA GPU-ovima.

Prilagođeni modeli, koje su izgradili različiti timovi u kompaniji Perplexity, hostuju se interno na NVIDIA Hopper i Blackwell GPU-ovima korišćenjem ROSE mehanizma za inferenciju koji je izgradio tim za AI inferenciju. ROSE je prvobitno napravljen radi bolje usluge prilagođenih Llama modela za dekodiranje jezika i klasifikaciju, sa interfejsom kompatibilnim sa NVIDIA Triton Inference Server-om. Vremenom je izrastao u potpuno razvijeni mehanizam za inferenciju koji pruža značajne LLM-ove, kao i rangiranje, klasifikaciju, bodovanje i druge modele zasnovane na transformatorima za širok spektar timova koji pokreću proizvode kompanije Perplexity. Dizajnerski cilj ROSE-a je da bude jednostavan: vremenom se ulaže napor da se dostignu maksimalne performanse, ali glavni prioritet je prilagođavanje arhitekturama i slučajevima upotrebe koji se stalno razvijaju za naše modele.

Prostor konfiguracije kernela kroz tipove podataka, dimenzije glava, konstante u vreme kompajliranja i postavke pokretanja.

Na visokom nivou, ROSE se sastoji od mehanizma koji prilagođava model interfejsu sa kojim klijenti mogu da komuniciraju. Mehanizmi se mogu u velikoj meri ponovo koristiti, upravljajući raspoređivanjem zahteva i grupisanjem (batching), kao i preuzimanjem dekodiranih tokena, vektorskih prikaza ili rezultata iz osnovnih modela. Inicijalizacija uređaja, komunikacioni kanali između procesa i čvorova, kao i učitavanje težina, su u njihovoj nadležnosti. Za LLM-ove, mehanizam je odgovoran za grupisanje (batching), deljenje na delove (chunking) i uzorkovanje. Takođe je odgovoran za alokaciju KV memorije, kako za punu tako i za linearnu pažnju, kao i za podudaranje prefiksa. Mehanizmi za vektorske prikaze (embeddings) su jednostavniji jer obavljaju samo grupisanje u realnom vremenu. Izvršavanje je optimizovano tako da mehanizam priprema sledeću rundu rada dok model radi paralelno na akceleratoru.

ROSE izlaže bogatu kolekciju prilagođenih slojeva za modeliranje, koji obavijaju GPU kernele izgrađene pomoću CuTeDSL-a, Tritona, CUDA-e, CUTLASS-a i cuBLAS-a. Za implementacije množenja matrica i pažnje, prvenstveno se oslanjamo na NVIDIA kernele iz CUTLASS-a i cuBLAS-a. Preostali deo cevovoda (pipeline), uključujući kernele za vektorske prikaze, normu, MoE i aktivaciju, čini deo biblioteke kernela ROSE-a. Kako bismo najbolje podržali hardver na kojem vršimo inferenciju, ROSE je opremljen sopstvenim primitivima komunikacije između uređaja za MoE rutiranje, otpremanje i kombinovanje. Iako su ostali kerneli relativno jednostavni, oni moraju biti temeljno specijalizovani kako bi se postigle maksimalne performanse na svim različitim skrivenim dimenzijama i uređajima na kojima rade. Pored toga, oslanjamo se na specijalizaciju u vreme kompajliranja kako bismo prilagodili sve opcije i podešavanja različitih modela (pristrasnost težina u odnosu na odsustvo pristrasnosti težina za RMS normu, funkcije aktivacije itd.). Shodno tome, odlučili smo da koristimo DSL-ove za njihovu implementaciju: prvobitno smo izgradili ROSE koristeći Triton. Međutim, u međuvremenu smo usvojili CuTeDSL kao primarno okruženje za GPU programiranje.

Sa matematičke perspektive, kerneli za inferenciju su obično jednostavni. Shodno tome, prilikom razvoja kernela, složenost se generalno ne nalazi u izražavanju toga koje izračunavanje treba da se obavi, već u tome kako tačno hardver treba to da uradi. Da bismo izgradili najbolje kernele, potreban nam je jezik koji nudi prave apstrakcije nad osnovnim hardverom. Pravi jezik nam omogućava da definišemo operacije na visokom nivou za izvršavanje izračunavanja, agresivno ih specijalizujući i optimizujući kroz sve tačke u širokom prostornom opsegu.

CuTeDSL je DSL zasnovan na Python-u za NVIDIA GPU-ove izgrađen na vrhu CuTe algebre rasporeda (layout algebra) i MLIR-a. On se kompajlira neposredno pre izvršenja (just-in-time) u visoko optimizovani PTX, uz zadržavanje kontrole nad hardverskim primitivima niskog nivoa. Cilj mu je da ponudi iste funkcije i performanse kao CUTLASS C++, bez skupog vremena kompajliranja koje koči brzinu razvoja. Od jezika koje smo istraživali ili sa kojima smo radili (uključujući CUDA, cuTile, TVM IR i Triton), smatramo da najviše odgovara razvoju naših kernela.

Sve u svemu, iako nije toliko visokog nivoa i sažet kao Triton ili cuTile, otkrili smo da su CuTe algebra rasporeda i njene kompozicije (izložene preko cute.Tensor) predstavljale razumne zamene. Pored toga, CuTeDSL nudi neograničen pristup hardverskim primitivima, obezbeđujući nam nivo kontrole koji nam je potreban za postizanje maksimalnih performansi u svim graničnim slučajevima. Sa CUDA-om, mogućnosti za specijalizaciju u vreme kompajliranja bile su ograničene, pošto je svaki dodatni parametar dodavao neprihvatljiv trošak vremenu kompajliranja ili složenost sistemu za građenje. Nasuprot tome, sa DSL-om koji se kompajlira neposredno pre izvršenja, možemo agresivnije da se specijalizujemo bez žrtvovanja brzine iteracije. Za razliku od TVM IR-a, koji je koristio CUDA-u kao međukorak, otkrili smo da je cevovod zasnovan na MLIR-u koji se spušta na PTX robusniji i lakši za otklanjanje grešaka od onog koji emituje CUDA-u kao međukorak. Pored toga, uz upotrebu prilagođenih operacija za obavijanje PTX instrukcija, otkrili smo da možemo u potpunosti da iskoristimo sve funkcije osnovnog hardvera.

Glavni deo troškova kompajliranja CUDA kernela obično se troši u fazi širenja šablona (template expansion) kernela, pošto se optimizacija i kompajliranje od PTX faze nadalje mogu završiti za nekoliko sekundi. Sa mnogo parametara u vreme kompajliranja, ova faza kompajliranja može da se produži na više minuta za netrivijalan kernel. Ovo usko grlo je naglašeno činjenicom da se prilike za paralelno kompajliranje javljaju tek kasnije, na nivoima IR i PTX, pri čemu se rukovanje C++ šablonima odvija sekvencijalno. Iako postoje trikovi za rešavanje ovoga, pomeranjem instanciranja šablona u odvojene jedinice kompajliranja, oni zahtevaju značajna prilagođavanja sistema za građenje i upotrebu glomaznih makroa. Nasuprot tome, CuTeDSL kerneli koji se kompajliraju neposredno pre izvršenja (JIT) ne zahtevaju složena pravila građenja i ne koče vreme kompajliranja, jer se razvojni ciklus može fokusirati na jednu tačku u prostoru konfiguracije, dok pažljivi unit testovi kasnije mogu pokriti celokupan prostor dizajna kernela.

Iako JIT kompajliranje prebacuje troškove sa vremena razvoja na vreme izvršavanja, režijski troškovi su više nego prihvatljivi. Kada se model primeni, on aktivira samo nekoliko tačaka iz celokupnog prostora konfiguracije, što izaziva režijske troškove u sekundi tokom pokretanja/zagrevanja mehanizma za inferenciju. Nasuprot tome, predkompajliranje svih konfiguracija kernela može da traje više minuta. Pored toga, režijski troškovi se mogu dodatno ublažiti keširanjem JIT kompajliranih artefakata.

Teški troškovi kompajliranja CUDA-e naglašeni su težinom praćenja porekla grešaka kada se koristi otpremanje i specijalizacija zasnovana na šablonima, kao što je slučaj sa CuTe-om i CUTLASS-om. Sa DSL-om, greške se mogu usmeriti na liniju koda koja ih je generisala, što značajno olakšava otklanjanje grešaka. Pored toga, mnogo je jednostavnije mapirati generisani kod na DSL izraz koji ga je generisao preko srednjih MLIR faza ili oslanjanjem na informacije za otklanjanje grešaka priložene uz generisani PTX. Zahvaljujući ovim aspektima DSL-a, otkrili smo da smo mnogo produktivniji pri razvoju kernela.

Iako kerneli dele mnoge sličnosti među modelima, oni i dalje moraju biti opremljeni nizom konfiguracionih opcija koje mogu uticati na performanse. Relevantni parametri uključuju:

Skrivena dimenzija / Dimenzija glave / Broj glava pažnje

Tipovi ulaznih/izlaznih podataka

Šema kvantizacije

Konstante u vreme kompajliranja (pristrasnosti, funkcije aktivacije itd.)

ROSE arhitekturu, uključujući njen engine, model, biblioteku kernela i servise.

Specijalizacija u vreme kompajliranja je korisna i za performanse i za otklanjanje grešaka. Posebno u cevovodnom procesoru, kao što je slučaj sa GPU-ovima, nepostojeća grana je bolja od grane koja se uvek ili nikada ne uzima. Pored toga, kompajliranje kernela za određeni skup konfiguracionih opcija rezultira manjim brojem PTX/SASS instrukcija, koje je lakše analizirati i za koje je lakše otklanjati greške kada se koristi cuda-gdb tokom istrage.

Za datu veličinu serije (batch size) i konfiguraciju, mogu se proceniti očekivani saobraćaj ka DRAM-u i broj aritmetičkih operacija. Nakon toga, optimalne odluke se mogu doneti od slučaja do slučaja za optimalnu alokaciju mreže (grid allocation), sinhronizaciju i upotrebu određenih hardverskih funkcija. CuTeDSL nam omogućava da definišemo i ponovo koristimo komponente za duboku specijalizaciju bez stvaranja neprihvatljivih režijskih troškova u vreme kompajliranja.

Generalno, performanse kernela za inferenciju moraju biti optimizovane za dva odvojena slučaja upotrebe:

Prefill: mala veličina serije (malo delova, sekvenci), mnogo tokena po sekvenci. Propusni opseg treba da bude maksimalan i za vektorske prikaze i za ulazne tokene LLM-ova. Latencija u opsegu od desetina do stotina milisekundi.

Decode: veličine serija od malih do velikih, jedan (bez spekulativnog dekodiranja) do nekoliko tokena po sekvenci (spekulacija ciljnog nacrta, predviđanje više tokena). Latencija blizu latencije pokretanja kernela.

Za GEMM i pažnju, faza predpopunjavanja (prefill) radi na velikom broju tokena i performanse su ograničene propusnim opsegom izračunavanja uređaja. Dekodiranje uključuje istu operaciju na malim veličinama serija, pri čemu u vremenima izvršavanja dominira trošak učitavanja aktivacija i težina iz DRAM-a u striming multiprocesore i tenzorske jezgre, što je umesto toga ograničeno memorijskom propusnošću uređaja. Iako je veličina serije ili ukupan broj tokena u seriji relativno jasan pokazatelj, prag prebacivanja između predpopunjavanja i dekodiranja je suptilniji jer je memorijska propusnost takođe određena skrivenom dimenzijom modela i tipom podataka aktivacije. Kerneli su posljedično podijeljeni na implementacije predpopunjavanja i dekodiranja ako bi algoritamske razlike ili dodjele mreže/bloka davale prednost bilo kome od njih.

Parametar koji je univerzalno različit između kernela za predpopunjavanje i dekodiranje jeste broj warpova u bloku (broj niti, podeljen sa veličinom warp-a od 32 niti po warp-u). Poželjno je da ovaj parametar bude konstanta u vreme kompajliranja, jer može pomoći u alokaciji registara. U CUDA-i se to obično radi navođenjem broja niti preko __launch_bounds__ i obezbeđivanjem broja warp-ova kao parametra šablona:

cpp

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

}

Broj registara dostupnih svim nitima u bloku je hardversko ograničenje, 64K na Hopper-u i Blackwell-u. Ako je broj niti poznat unapred, kompajler tačno zna koliko registara može da dodeli svakoj niti i može uspešno da kompajlira agresivnije optimizovane kernele. Međutim, ovo dolazi sa cenom u vremenu kompajliranja, pošto se sada kernel mora ponovo kompajlirati za različite veličine warp-ova. Sa CuTeDSL-om ovo više nije problem, pošto ne moramo unapred da plaćamo cenu kompajliranja za celokupan prostor konfiguracije. DSL takođe očekuje da veličina bloka bude konstantan izraz, implicitno vršeći optimizaciju koju sugeriše __latch_bounds__ u CUDA-i.

Pored toga što broj warp-ova čini konstantom, poželjno je da se on razlikuje između predpopunjavanja i dekodiranja. Da bi se u potpunosti iskoristila memorijska propusnost GPU-a, nije dovoljno uvek dodeliti veliki broj warp-ova po bloku, pošto se puna propusnost može postići samo korišćenjem što većeg broja dostupnih striming multiprocesora (SM). Shodno tome, za dati broj tokena i poznate tipove podataka i skrivene dimenzije, prvo biramo idealnu dimenziju vektora po niti (8 za bfloat16, 4 za float32) kako bismo osigurali da generisani kod koristi 128-bitne instrukcije učitavanja/skladištenja (256-bitne na Blackwell-u). Znajući veličinu vektora, možemo izračunati broj warp-ova potrebnih kernelu. Prvo pokušavamo da maksimizujemo broj SM-ova: ako ima manje warp-ova nego SM-ova, bira se jedan warp po bloku. Ovo je obično varijanta dekodiranja kernela. Ako ima dovoljno warp-ova, postepeno povećavamo broj warp-ova na 8, 16 ili 32, vršeći benchmark gornje granice nakon koje obično više ne vidimo poboljšanje u performansama.

Iako je Triton obavljao neke od ovih optimizacija ispod haube, otkrili smo da možemo postići slične rezultate pažljivim korišćenjem CuTe algebre rasporeda, bez ograničenja od strane nedostataka kompajlatora visokog nivoa. Naročito kada su uključene horizontalne redukcije unutar pod-warp-a, otkrili smo da nam je CuTeDSL ponudio finu kontrolu potrebnu za generisanje optimalnog koda.

U nekim situacijama, da bi se rad rasporedio na što više blokova, cooperative_groups::this_grid().sync() se može koristiti za sinhronizaciju svih blokova koji učestvuju. Važe ograničenja: pošto se blok ne može ukloniti sa rasporeda sa SM-a jednom kada započne izvršavanje, da bi se garantovao napredak unapred, broj blokova ne može premašiti broj SM-ova dostupnih na uređaju, kako bi se osiguralo da u bilo kom trenutku svi blokovi mogu da ažuriraju i anketiraju lokacije globalne memorije uključene u sinhronizaciju. Ovo ograničenje nameće cudaLaunchCooperativeKernel. Cena mrežne barijere je oko 3us na Hopper-u, plus cena pražnjenja svih prethodnih memorijskih operacija preko implicitne memorijske barijere u nekim situacijama. Sve u svemu, ovo je uporedivo sa režijskim troškovima deljenja kernela na dva pokretanja, što vodi do donošenja odluke o dizajnu osetljive na performanse. Uprkos činjenici da Programski zavisno pokretanje (PDL) može smanjiti režijske troškove pokretanja počevši od Hopper-a, za dekoderske kernele izbegavanje ovih režijskih troškova se isplati.

Ako jedno od opterećenja zahteva više blokova od broja SM-ova, kernel treba podeliti, oslanjajući se na implicitnu barijeru pokretanja kernela. Ovo je prvenstveno slučaj za predpopunjavanje. Ako broj blokova koji je idealan za sva opterećenja staje u mrežu, preferira se mrežna barijera, kao što je slučaj sa kernelima za dekodiranje koji rade na manjem broju tokena.

U CuTeDSL-u, ove odluke se mogu doneti u kodu za pokretanje na strani domaćina (host-side) i kerneli se mogu sastavljati i preplitati sa mrežnim barijerama po potrebi, kao što je ilustrovano u primeru ispod. Sa takvom hibridnom strategijom, kerneli za dekodiranje mogu biti optimizovani da dostignu minimalne latencije, a oni za predpopunjavanje mogu biti optimizovani za maksimalni propusni opseg. Za razliku od rešenja zasnovanih na CUDA-i, ova logika se može definisati sažetije i uz manje režijskih troškova u vreme kompajliranja.

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)
		)

U slučaju naših MoE all-to-all kernela, kao što je detaljno opisano u nastavku, posedovanje fleksibilnosti za izgradnju kernela sa mrežnim barijerama ili bez njih i njihova optimizacija u vreme kompajliranja pomogli su nam da ostvarimo značajna poboljšanja performansi kroz CuTeDSL.

Prvi kernel koji smo izgradili u CuTeDSL-u bio je specijalizovani RMS norm kernel, optimizovan za male dimenzije. QK norma, koja se koristi u sloju samopažnje nedavnih Qwen modela, primenjuje se po glavi (per-head), na 64, 128 ili 256 elemenata. Za razliku od regularne RMS norme, koja radi na skrivenim dimenzijama koje premašuju hiljade, zahtevajući horizontalnu redukciju kroz deljenu memoriju i registre, QK norma zahteva redukciju samo preko warp-ova, uz mogućnost sabijanja više glava u jedan warp. Ovo je postao slučaj upotrebe gde smo mogli agresivno da optimizujemo u vreme kompajliranja da bismo pronašli najbolje dodele tokena u warp za sve dimenzije glava i tipove podataka.

RMS norma je prilično jednostavna operacija, mada zahteva horizontalnu sumu za izračunavanje kvadratne norme vektora:

cpp

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

Pošto su glave pažnje obično male, a operacije po niti u velikoj meri vektorizovane, jedan warp može izračunati kvadratnu normu. Da bi se maksimalno iskoristila memorijska propusnost, svaka nit treba da obradi 128 bita ulaza, što se pretvara u 8 bfloat16 vrednosti. Shodno tome, polovina warp-a je dovoljna za obradu glave, a kvadratna norma se može odrediti korišćenjem horizontalne redukcije koju obavlja shuffle_sync_bfly. Počevši od dimenzije glave i tipa podataka, možemo odrediti broj glava po warp-u i navesti argumente i broj rundi redukcije u vreme kompajliranja.

Tako možemo izgraditi specijalizovani kernel sa relativno malo linija koda, oslanjajući se na automatski vektorizovana učitavanja/skladištenja i koristeći operacije nižeg nivoa specifične za arhitekturu za redukciju warp-a. Kernel određuje na kojoj glavi (Q ili K) treba da radi i identifikuje odgovarajuće težine koje treba primeniti. On učitava ceo vektor iz glave i prvo ga lokalno sumira. Naknadno, na osnovu broja niti po glavi, izračunava se kvadratna norma. Na kraju, RMS norma, automatski vektorizovana, se određuje i upisuje nazad u memoriju. Parametri specifični za model, kao što su WEIGHT_BIAS i EPS, su svedeni na konstante (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)

U poređenju sa regularnim RMS norm kernelom, ova specijalizovana varijanta pruža 2-3 puta bolje performanse u pogledu propusnog opsega i latencije i za predpopunjavanje i za dekodiranje.

MoE otpremanje i kombinovanje mešaju ulazne tokene do stručnjaka kojima su usmereni, obavljajući ponderisanu akumulaciju nakon što MoE GEMM kerneli obave relevantna množenja matrica. Performanse su od posebnog značaja, jer se različiti proizvodi oslanjaju na MoE modele različitih veličina. Na osnovu tipova primene, razlikujemo različite kernele za optimizaciju:

TP=1: Lkalni kernel za jedan GPU. Otpremanje meša tokene do relevantnih stručnjaka, a kombinovanje ih akumulira. Koristi se prvenstveno za Qwen3 30B primene na H100.

DP=1, TP ≠ 1: Tenzorski paralelni kernel za jedan čvor (singlenode). Ulaz se replikuje, otpremanje bira tokene za trenutni rang, a kombinovanje ih akumulira kroz sve rangove. Akumulacija preko NVLink-a ili NVLS-a. Koristi se za velike sisteme za predpopunjavanje (prefillers).

DP ≠ 1, NVLink: Podatkovno paralelni kernel sa opcionom tenzorskom paralelnošću. Svaki rang šalje tokene svakom drugom rangu koji hostuje određene stručnjake. Informacije o rutanju se agregiraju, tokeni se prenose preko NVIDIA NVLink-a. Koristi se za podatkovno paralelne H200 ili GB200 primene.

DP ≠ 1, InfiniBand: Peer-to-peer podatkovno paralelni i tenzorski paralelni kernel koji radi preko i NVLink-a i NVIDIA Quantum InfiniBand-a. Koristi se za višecvorne Hopper primene.

Lokalni kernel broji broj tokena dodeljenih svakom ekspertu na osnovu tabele rutiranja i dodeljuje jedinstveni indeks svakom tokenu unutar liste tokena koji pripadaju svakom ekspertu. Nakon koraka rutiranja, tokeni se kopiraju u opciono dopunjene baferе kako bi se prosledili MoE GEMM kernelu. Koristeći lokacije tokena, kernel za spajanje zatim vrši akumulaciju u registrima, upisujući završne tokene u izlazni bafer.

Faza otpremanja varira između predpopunjavanja i dekodiranja: za malo tokena ima malo informacija o rutanju za čitanje, tako da ih obrađujemo u jednom bloku i signaliziramo ostatak mreže, pri čemu pojedinačni blokovi kopiraju delove tokena. Dodela tokena se određuje atomskim uvećavanjem brojača u deljenoj memoriji. Pošto predpopunjavanje mora da obradi više informacija o rutanju, rutiranje je podeljeno na više blokova u zasebnom kernelu. Indeksi se prvo dodeljuju lokalno pomoću brojača u deljenoj memoriji, a zatim se agregiraju kroz sve blokove u globalnoj memoriji da bi se odredila pozicija svakog tokena. Mreža ovog kernela je ograničena na broj SM-ova kako bi se mogle koristiti mrežne barijere. Naknadni kernel obavlja kopiranja, pokrenut preko mnogo većeg broja blokova i warp-ova kako bi se u potpunosti iskoristio dostupan memorijski propusni opseg.

Kernel za kombinovanje je relativno jednostavan jer obavlja vektorizovanu akumulaciju, oslanjajući se na informacije o rutanju i pomeraje (offsets) koje izračunava otpremanje. CuTeDSL se ovde pokazao dragocenim, pošto se izračunavanje moglo izraziti kroz jednostavne CuTe tenzore koji se kompajliraju u efikasan kod za sve tipove podataka i skrivene dimenzije.

U tenzorski paralelnoj primeni, informacije o rutanju i ulazni tokeni se replikuju. Lokalni kernel za otpremanje se može u potpunosti ponovo koristiti uz malu promenu za rutiranje i izbor samo onih tokena dodeljenih stručnjacima koji se nalaze na trenutnom rang-u. Uz neke parametre i specijalizaciju u vreme kompajliranja, lokalni kerneli za otpremanje se ponovo koriste.

Pošto je broj rangova obično manji od broja stručnjaka kojima je token usmeren, u prosečnom slučaju se može pretpostaviti da svaki token ima bar jednu od svojih replika usmerenu na svaki rang. U ovoj situaciji, lokalni kernel za kombinovanje se može ponovo koristiti za agregiranje lokalnih kopija tokena, nakon čega sledi all-reduce za sabiranje kroz rangove koji učestvuju. U slučaju kernela za predpopunjavanje, pokrećemo odvojene kernele, međutim dekodiranje spaja obe operacije u jedno pokretanje.

Za all-reduce koristimo običan NVLink ako učestvuju samo dva ranga, prebacujući se na NVLS kada je uključeno 4 ili više rangova radi povećanja propusnog opsega od 10-20%. Zahvaljujući DSL-u, možemo deliti većinu logike redukcije i razmene vršnih baifera između MoE i regularnih all-reduce kernela. Ovo je situacija u kojoj je bliska kontrola nad hardverom korisna, pošto su DSL operacije nad multimem.ld_reduce i multimem.st neophodne za NVLS. Sama implementacija redukcije u velikoj meri prati all_reduce primer iz NCCL testova.

Kerneli paralelni sa podacima oslanjaju se na NVLink za agregiranje tokena za više uređaja. U fazi otpremanja, pre razmene tokena, vršnjaci (peers) se prvo lokalno usmeravaju da identifikuju ciljni pomeraj na svakom rang-u za svaki token koji šalju. Naknadno, počinju da kopiraju tokene u prijemne baifere dodeljene za svakog pošiljaoca kako bi se izbegli sukobi. Upisi sa otpuštanjem i skladištenjem po tokenu (per-token release-store writes) koriste se za signaliziranje završetka. Na strani prijemnika, informacije o rutanju se agregiraju kako bi se svi tokeni mogli sažeti u jedan baifer za MoE GEMM. Prijemnici anketiraju fleg koji su pisali vršnjaci sa acquire-load-om da sačekaju da se upisi završe.

Na strani kombinovanja, tokeni se kopiraju preko NVLink-a do izvornih rangova, sa flegom po tokenu koji označava prijem. Kernel za kombinovanje je sličan lokalnoj i singlenode varijanti, sa malim razlikama u adresiranju i nekom dodatnom logikom sinhronizacije.

Izgradili smo biblioteku PTX omotača (wrappers) da izložimo release-acquire učitavanja i skladištenja opsega na nivou sistema i uređaja za implementaciju ovog kernela. Poboljšanjem granularnosti flegova sinhronizacije i alokacijom više blokova za slanje i primanje tokena preko NVLink-a, ostvarili smo dvostruka poboljšanja u odnosu na naše prethodne CUDA kernele.

Naši peer-to-peer kerneli su prethodno postigli performanse najvišeg nivoa (state-of-the-art) na Hopper GPU-ovima uparenim sa ConnectX-7 InfiniBand adapterima. U međuvremenu smo portovali kernele za kombinovanje na CuTeDSL radi daljeg povećanja performansi. Pošto su CUDA verzije već u velikoj meri se oslanjale na specijalizaciju u vreme kompajliranja, pokazale su se kao uska grla u vreme kompajliranja u našem repozitorijumu, primoravajući nas da izostavimo neke parametre. Migracijom na JIT-kompajlirani DSL, specijalizovali smo se za više parametara i prilagodili alokacije mreže, uvodeći dalja razilaženja između predpopunjavanja i dekodiranja za oko 3% ubrzanja latencije.

Iz perspektive ROSE kernela, podrška za Blackwell spada u dve kategorije: agnostički kerneli u odnosu na arhitekturu i kerneli specifični za arhitekturu. Kerneli specifični za arhitekturu, oslanjajući se na tenzorske jezgre, zahtevaju eksplicitan rad radi prilagođavanja. Međutim, zahvaljujući upotrebi DSL-a, kerneli agnostički u odnosu na arhitekturu se prilagođavaju bez muke.

Blackwell GPU ima do 160 SM-ova, za razliku od Hopperovih 132. Za većinu kernela ovo je konstanta u vreme kompajliranja, a logika u kodu pokretača na strani domaćina to uzima u obzir pri određivanju optimalnog kernela za pokretanje (predpopunjavanje u odnosu na dekodiranje). Pored toga, Blackwell takođe uvodi 256-bitne instrukcije učitavanja, koje se lako mogu emitovati povećanjem širine vektora (poslednja dimenzija CuTe tenzora) radi poboljšanja iskorišćenja memorijske propusnosti.

Blackwell uvodi hardversku podršku za značajan broj šema kvantizacije (MXFP8, MXFP4, NVFP4). Podrška za ove šeme je primenjena prvenstveno prilagođavanjem postojećih kernela za skaliranje po blokovima za FP8. JIT podrška je ovde ključna, pošto se novi režimi mogu podržati bez eksplozije prostora konfiguracije u vreme kompajliranja. U prošlosti smo otkrili da je Triton nepraktičan za izražavanje horizontalnih redukcija kroz nekoliko elemenata, preko polovine ili četvrtine warp-a. Finitija kontrola nad redukcijama bila je ključna za postizanje optimalnih performansi pri izračunavanju apsolutne maksimalne vrednosti preko 16 ili 32 elementa za NVFP4 i MXFP8/MXFP4, redom.

Kao relativno nov jezik, LLM-ovi su imali malo kontakta sa CuTeDSL kodom tokom svog predobučavanja. Pored toga, stil našeg sadašnjeg koda, sa kernelima optimizovanim za aktivacije i norme, bio je malo drugačiji od postojećih CUTLASS ili Quack primera, koji su bili u velikoj meri optimizovani za TMA i instrukcije množenja matrica specifične za arhitekturu. Dostići tačku u kojoj LLM-ovi mogu pomoći u pisanju kernela bio je postepen proces.

Ručno smo izgradili početni korpus kernela i pomoćnih funkcija, prateći dosledne šablone, pre nego što smo mogli da se oslonimo na agente za konverziju kernela. Kako je korpus napredovao, mogli smo da obavimo grube prevode sa CUDA-e na CuTeDSL, uz značajan naknadni ručni rad. Konkretno, otkrili smo da su se modeli mučili sa konverzijama pokazivača u tenzor i vektorizacijom, zaglavljujući se u pokušajima da isprave greške tipova kada su težili da nateraju test da se kompajlira. Pokazali su malo razumevanja prirode kernela zasnovane na mrežama, mešajući sinhronizaciju niti i mreže. Najviše frustrirajuće je bilo iskušenje da se haluciniraju nepostojeći prevodi biblioteka specifičnih za CUDA-u umesto građenja na našim sopstvenim CuTeDSL zamenama. Vremenom, kako je korpus CuTeDSL kernela rastao, otkrili smo da LLM-ovi mogu da izgrade većinu skeletnog koda i popune praznine, omogućavajući nam da se fokusiramo na dizajn visokog nivoa.

U poređenju sa CUDA-om, smatramo da CuTeDSL pruža mnogo bolje performanse u vreme kompajliranja, omogućavajući nam da se specijalizujemo za više parametara i postignemo poboljšanja performansi. Iako su kerneli malo opsirniji zbog nedostatka određenih apstrakcija u poređenju sa Tritonom, smatramo da je složenost opravdana jer nam finija kontrola nad hardverom omogućava da izgradimo brže kernele.

U vreme pisanja ovog teksta, jaz u performansama između CUTLASS C++ i CuTeDSL kernela je minimalan, što otvara uzbudljive prilike za sastavljanje i spajanje GEMM kernela visokih performansi sa komunikacijom između uređaja, normalizacijom ili aktivacijama. Posvećivanjem ovom DSL-u, moći ćemo da izvučemo dodatne performanse iz naših uređaja kako bismo poboljšali propusni opseg kroz naša opterećenja, od vektorskih prikaza do LLM-ova.

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}
}

Inferencija u kompaniji Perplexity AI

Zašto CuTeDSL?

Ispustvo za programere

Veliki prostor konfiguracije

Unapred popuni i dekodiraj specijalizaciju

Sinhronizacija mreže (Grid synchronization)

CuTeDSL kerneli

QK norma

Mešavina stručnjaka (Mixture-of-Experts) Otpremanje/Kombinovanje (Dispatch/Combine)

Lokalno

Jedan čvor (Singlenode)

Paralelno sa podacima (Data Parallel)

Peer-to-Peer

Podrška za Blackwell

Generisanje kernela potpomognuto veštačkom inteligencijom

Zaključak i budući rad

Reference

Citiranje