CuTeDSL u tvrtki Perplexity

U tvrtki Perplexity koristimo vlastiti uslužni mehanizam optimiziran za izvođenje (ROSE) za posluživanje modela koji se kreću od ugrađivanja do jezičnih modela s bilijunima parametara (LLM). Svi API-ji, uključujući Sonar, pretraživanje i ugrađivanja, oslanjaju se na naš mehanizam koji je osmišljen kako bi omogućio istraživačima istraživanje najprikladnijih modela i arhitektura, uz istovremeno brzo uvođenje najboljih kandidata u produkciju. Kako bismo to postigli, moramo biti u mogućnosti brzo implementirati, uvesti i optimizirati zaključivanje za najsuvremenije arhitekture. U ovom blogu istražujemo kako nam jezgre CuTeDSL integrirane u naš mehanizam za zaključivanje omogućuju izgradnju potrebnih GPU jezgri kako bismo dosegnuli vrhunske performanse modela na NVIDIA GPU-ovima.

Prilagođeni modeli, koje su izgradili različiti timovi u tvrtki Perplexity, hostaju se u vlastitom okruženju na NVIDIA Hopper i Blackwell GPU-ovima pomoću mehanizma za zaključivanje ROSE koji je izgradio tim za AI zaključivanje. ROSE je izvorno izgrađen za bolje posluživanje prilagođenih modela Llama za dekodiranje jezika i klasifikaciju, s sučeljem kompatibilnim s poslužiteljem za zaključivanje NVIDIA Triton. S vremenom je izrastao u punopravni mehanizam za zaključivanje koji poslužuje značajne LLM-ove, kao i rangiranje, klasifikaciju, bodovanje i druge modele temeljene na transformatorima za širok raspon timova koji pokreću Perplexityjeve proizvode. Cilj dizajna mehanizma ROSE jest jednostavnost: s vremenom se ulaže trud kako bi se dosegla vrhunska performansa, no glavni prioritet je prilagodba stalno promjenjivim arhitekturama i slučajevima korištenja za naše modele.

Prostor konfiguracije kernela preko tipova podataka, dimenzija glava, konstanti u vremenu prevođenja i postavki pokretanja.

Na visokoj razini, ROSE se sastoji od mehanizma koji prilagođava model sučelju s kojim klijenti mogu komunicirati. Mehanizmi su visoko višekratno iskoristivi, upravljaju raspoređivanjem i grupiranjem zahtjeva te dohvaćaju dekodirane tokene, ugrađivanja ili bodove iz temeljnih modela. Njima se upravlja inicijalizacijom uređaja, komunikacijskim kanalima između procesa i čvorova te učitavanjem težina. Za LLM-ove mehanizam je odgovoran za grupiranje, dijeljenje u blokove i uzorkovanje. Također je odgovoran za alokaciju KV pohrane, kako za potpunu tako i za linearnu pažnju, kao i za podudaranje prefiksa. Mehanizmi za ugrađivanje su jednostavniji jer obavljaju samo internetsko grupiranje. Izvršavanje je optimizirano tako da mehanizam priprema sljedeći krug rada dok model paralelno radi na akceleratoru.

ROSE izlaže bogatu zbirku prilagođenih slojeva za modeliranje, omatajući GPU jezgre izgrađene s pomoću CuTeDSL, Triton, CUDA, CUTLASS i cuBLAS. Za implementacije množenja matrica i pažnje prvenstveno se oslanjamo na NVIDIA jezgre iz CUTLASS-a i cuBLAS-a. Ostatak cjevovoda, uključujući jezgre za ugrađivanje, normu, MoE i aktivaciju, dio je knjižnice jezgri mehanizma ROSE. Kako bi se što bolje podržao hardver na kojem pokrećemo zaključivanje, ROSE je opremljen vlastitim primitivima komunikacije između uređaja za usmjeravanje, slanje i kombiniranje MoE-a. Iako su druge jezgre relativno jednostavne, moraju biti temeljito specijalizirane kako bi se postigle vrhunske performanse na svim različitim skrivenim dimenzijama i uređajima na kojima se izvode. Osim toga, oslanjamo se na specijalizaciju u vrijeme kompilacije kako bismo udovoljili svim postavkama i modifikacijama različitih modela (pristranost težine u odnosu na nedostatak pristranosti težine za RMS normu, aktivacijske funkcije itd.). Posljedično, odlučili smo koristiti DSL-ove za njihovu implementaciju: izvorno smo izgradili ROSE koristeći Triton. Međutim, u međuvremenu smo usvojili CuTeDSL kao primarno okruženje za GPU programiranje.

Iz matematičke perspektive, jezgre za zaključivanje obično su jednostavne. Shodno tome, prilikom razvoja jezgre složenost općenito ne leži u izražavanju koje se računanje treba obaviti, već u tome kako točno hardver to treba učiniti. Kako bismo izgradili najbolje jezgre, potreban nam je jezik koji nudi prave apstrakcije nad temeljnim hardverom. Pravi jezik omogućuje nam specificiranje operacija visoke razine za obavljanje računanja, uz agresivnu specijalizaciju i optimizaciju kroz sve točke u širokom prostornom konfiguracijskom prostoru.

CuTeDSL je DSL temeljen na Pythonu za NVIDIA GPU-ove izgrađen na vrhu algebre rasporeda CuTe i MLIR-a. Kompajlira se just-in-time u visoko optimizirani PTX, istovremeno zadržavajući kontrolu nad hardverskim primitivima niske razine. Cilj mu je ponuditi iste značajke i performanse kao CUTLASS C++, bez dugih vremena kompilacije koja koče brzinu razvoja. Od jezika koje smo pregledali ili s kojima smo radili (uključujući CUDA, cuTile, TVM IR i Triton), smatramo ga najboljim izborom za razvoj naših jezgri.

Sve u svemu, iako nije tako visoke razine i sažet kao Triton ili cuTile, otkrili smo da su algebra rasporeda CuTe i njezine kompozicije (izložene putem cute.Tensor) razumni zamjenski elementi. Osim toga, CuTeDSL nudi neograničen pristup hardverskim primitivima, pružajući nam razinu kontrole koja nam je potrebna za postizanje vrhunskih performansi u svim rubnim slučajevima. S jezikom CUDA mogućnosti za specijalizaciju u vrijeme kompilacije bile su ograničene, jer je svaki dodatni parametar dodavao nerazuman trošak vremenima kompilacije ili složenost sustavu građenja. Suprotno tome, s DSL-om kompiliranim u trenutku izvršavanja (Just-In-Time Compiled), možemo se agresivnije specijalizirati bez žrtvovanja brzine iteracije. Za razliku od TVM IR-a, koji je koristio CUDA-u kao međukorak, otkrili smo da je cjevovod temeljen na MLIR-u koji se spušta na PTX robusniji i lakši za otklanjanje pogrešaka od onoga koji je emitirao CUDA-u kao međukorak. Dodatno, uz korištenje prilagođenih operacija za omatanje PTX instrukcija, otkrili smo da možemo u potpunosti iskoristiti sve značajke temeljnog hardvera.

Glavni dio troška kompilacije CUDA jezgri obično se troši u fazi širenja predložaka jezgri, budući da se optimizacija i kompilacija od PTX faze nadalje mogu završiti u roku od nekoliko sekundi. S mnogo parametara u vrijeme kompilacije, ova faza kompilacije može narasti na više minuta za netrivijalnu jezgru. Ovo usko grlo naglašeno je činjenicom da se prilike za paralelnu kompilaciju javljaju tek kasnije, na razinama IR-a i PTX-a, pri čemu se rukovanje C++ predlošcima odvija sekvencijalno. Iako postoje trikzi za rješavanje ovoga premještanjem instanciranja predložaka u zasebne jedinice kompilacije, oni zahtijevaju značajne prilagodbe sustava građenja i upotrebu nespretnih makronaredbi. Suprotno tome, CuTeDSL jezgre koje se kompiliraju u trenutku izvršavanja (JIT) ne zahtijevaju složena pravila građenja niti ometaju vremena kompilacije, budući da se razvojni ciklusi mogu usredotočiti na jednu točku u konfiguracijskom prostoru, dok kasniva pomna testiranja jedinica mogu pokriti cijeli projektni prostor jezgre.

Iako JIT kompilacija premješta troškove s vremena razvoja na vrijeme izvođenja, preopterećenje je više nego prihvatljivo. Kada se model uvede, on aktivira samo nekoliko točaka iz cijelog konfiguracijskog prostora, uzrokujući preopterećenja u trajanju od nekoliko sekundi tijekom pokretanja / zagrijavanja mehanizma za zaključivanje. Suprotno tome, predkompajliranje svih konfiguracija jezgre moglo bi trajati više minuta. Dodatno, preopterećenja se mogu dodatno ublažiti predmemoriranjem JIT-kompajliranih artefakata.

Teški trošak kompilacije za CUDA-u naglašen je poteškoćama u praćenju podrijetla pogrešaka kada se koriste raspačavanje i specijalizacija temeljeni na predlošcima, kao što je slučaj s CuTe i CUTLASS-om. S DSL-om se pogreške mogu usmjeriti na liniju koda koja ih je generirala, čime se značajno olakšava otklanjanje pogrešaka. Osim toga, puno je jednostavnije mapirati generirani kod na DSL izjavu iz koje potječe putem međufaza MLIR-a ili oslanjanjem na podatke o otklanjanju pogrešaka pridružene generiranom PTX-u. Zahvaljujući ovim aspektima DSL-a, otkrili smo da smo znatno produktivniji pri razvoju jezgri.

Iako jezgre dijele mnoge sličnosti među modelima, ipak moraju biti opremljene nizom konfiguracijskih gumbiju koji mogu utjecati na performanse. Relevantni parametri uključuju:

Skrivena dimenzija / dimenzija glave / broj glava pažnje

Ulazne/izlazne vrste podataka

Shema kvantizacije

Konstante u vrijeme kompilacije (pristranosti, aktivacijske funkcije itd.)

ROSE arhitektura, uključujući njezin engine, model, biblioteku kernela i usluge.

Specijalizacija u vrijeme kompilacije korisna je i za performanse i za otklanjanje pogrešaka. Posebno u cjevovodnom procesoru, kao što je slučaj s GPU-ovima, nepostojeća grana je bolja od grane koja se uvijek ili nikada ne uzima. Dodatno, kompilacija jezgre za određeni skup konfiguracijskih postavki rezultira s manje PTX/SASS instrukcija koje je lakše analizirati i za koje je lakše otkloniti pogreške kada se koristi cuda-gdb tijekom istrage.

Za zadanu veličinu serije i konfiguraciju mogu se procijeniti očekivani promet prema DRAM-u i broj aritmetičkih operacija. Nakon toga, optimalne odluke mogu se donositi od slučaja do slučaja za optimalnu alokaciju mreže, sinkronizaciju i korištenje određenih hardverskih značajki. CuTeDSL nam omogućuje definiranje i ponovno korištenje komponenti za duboku specijalizaciju bez izazivanja neprihvatljivog preopterećenja u vrijeme kompilacije.

Općenito, performanse jezgri za zaključivanje moraju biti optimizirane za dva odvojena slučaja korištenja:

Prefill: mala veličina serije (malo komada, sekvenci), mnogo tokena po sekvenci. Propusnost koju treba maksimalno povećati i za ugrađivanja i za ulazne tokene LLM-ova. Latencija u rasponu od desetaka do stotina milisekundi.

Dekodiranje: male do velike veličine serije, jedan (bez spekulativnog dekodiranja) do nekoliko tokena po sekvenci (spekulacija cilja nacrta, predviđanje više tokena). Latencija blizu latencije pokretanja jezgre.

Za GEMM i pažnju, faza ispune (prefill) izvodi se na velikom broju tokena i performanse su ograničene propusnošću računanja uređaja. Dekodiranje uključuje istu operaciju na malim veličinama serije, pri čemu vremenima izvršavanja dominira trošak učitavanja aktivacija i težina iz DRAM-a u streaming multiprocesore i tenzorske jezgre, što je umjesto 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 ispune i dekodiranja je suptilniji jer je memorijska propusnost također određena skrivenom dimenzijom modela i vrstom podataka aktivacije. Posljedično, jezgre se dijele na implementacije ispune i dekodiranja ako bi algoritamske razlike ili dodjele mreže/bloka davale prednost jednoj od njih.

Parametar koji je univerzalno drugačiji između jezgri za ispunu i dekodiranje jest broj warpova u bloku (broj niti, podijeljen s veličinom warpa od 32 niti po warpu). Poželjno je učiniti ovaj parametar konstantom u vrijeme kompilacije jer može pomoći u alokaciji registara. U CUDA-i se to obično radi specificiranjem broja niti putem __launch_bounds__ i davanjem broja warpova kao parametra predloška:

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 Hopperu i Blackwellu. Ako je broj niti poznat unaprijed, prevoditelj točno zna koliko registara može dodijeliti svakoj niti i može uspješno kompilirati agresivnije optimizirane jezgre. Međutim, to dolazi uz trošak kompilacije, jer se sada jezgra mora ponovno kompilirati za različite veličine warpova. S CuTeDSL-om to više nije problem, jer ne moramo unaprijed plaćati trošak kompilacije za cijeli konfiguracijski prostor. DSL također očekuje da veličina bloka bude konstantan izraz, čime se implicitno provodi optimizacija koju sugerira __latch_bounds__ u CUDA-i.

Osim što se broj warpova postavlja kao konstanta, poželjno je mijenjati ga između ispune i dekodiranja. Kako bi se u potpunosti iskoristila memorijska propusnost GPU-a, nije dovoljno uvijek dodijeliti visok broj warpova po bloku, budući da se puna propusnost može postići samo korištenjem što većeg broja dostupnih streaming multiprocesora (SM). Posljedično, za zadani broj tokena te poznate vrste podataka i skrivene dimenzije, najprije biramo idealnu vektorsku dimenziju po niti (8 za bfloat16, 4 za float32) kako bismo osigurali da generirani kod koristi 128-bitne instrukcije učitavanja/pohrane (256-bitne na Blackwellu). Poznavajući veličinu vektora, možemo izračunati broj warpova potrebnih za jezgru. Prvo pokušavamo maksimizirati broj SM-ova: ako ima manje warpova nego SM-ova, odabire se jedan warp po bloku. Ovo je obično varijanta jezgre za dekodiranje. Ako ima dovoljno warpova, postupno povećavamo broj warpova na 8, 16 ili 32, vršeći testiranje performansi gornje granice nakon koje obično više ne vidimo poboljšanje performansi.

Iako je Triton izveo neke od ovih optimizacija ispod haube, otkrili smo da možemo postići slične rezultate pažljivim korištenjem algebre rasporeda CuTe, a da nismo ograničeni nedostacima kompajlera visoke razine. Posebno kada su uključene horizontalne redukcije unutar podwarpa, otkrili smo da nam je CuTeDSL ponudio finu kontrolu potrebnu za generiranje optimalnog koda.

U nekim situacijama, kako bi se rad rasporedio na što više blokova, cooperative_groups::this_grid().sync() se može koristiti za sinkronizaciju svih sudjelujućih blokova. Primjenjuju se ograničenja: budući da se blok ne može rasporediti izvan SM-a nakon što započne izvršavanje, kako bi se zajamčio napredak prema naprijed, broj blokova ne smije premašiti broj SM-ova dostupnih na uređaju, kako bi se osiguralo da u bilo kojem trenutku svi blokovi mogu ažurirati i anketirati lokacije globalne memorije uključene u sinkronizaciju. Ovo ograničenje provodi cudaLaunchCooperativeKernel. Trošak mrežne barijere iznosi oko 3us na Hopperu, plus trošak pražnjenja svih prethodnih memorijskih operacija putem implicitne memorijske barijere u nekim situacijama. Sve u svemu, ovo je usporedivo s preopterećenjem dijeljenja jezgre u dva pokretanja, što dovodi do odluke o dizajnu koja je osjetljiva na performanse. Unatoč činjenici da programsko ovisno pokretanje (PDL) može smanjiti preopterećenje pokretanja od Hoppera nadalje, za jezgre dekodiranja isplati se izbjeći ovo preopterećenje.

Ako jedno od radnih opterećenja zahtijeva više blokova od broja SM-ova, jezgru bi trebalo podijeliti, oslanjajući se na implicitnu barijeru pokretanja jezgri. To je prvenstveno slučaj za ispunu. Ako broj blokova koji je idealan za sva radna opterećenja stane u mrežu, preferira se mrežna barijera, kao što je slučaj s jezgrama za dekodiranje koje rade na manjem broju tokena.

U CuTeDSL-u ove se odluke mogu donijeti u kodu pokretanja na strani hosta, a jezgre se mogu sastavljati i isprepletati s mrežnim barijerama prema potrebi, kao što je prikazano u primjeru u nastavku. S takvom hibridnom strategijom, jezgre za dekodiranje mogu se optimizirati da dosegnu minimalne latencije, a one za ispunu mogu se optimizirati za maksimalnu propusnost. Za razliku od rješenja temeljenih na CUDA-i, ova se logika može definirati sažetije i s manjim preopterećenjem u vrijeme kompilacije.

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 jezgri, kao što je detaljno opisano u nastavku, posjedovanje fleksibilnosti za izgradnju jezgri sa ili bez mrežnih barijera i njihova optimizacija u vrijeme kompilacije pomogle su nam ostvariti značajna poboljšanja performansi kroz CuTeDSL.

Prva jezgra koju smo izgradili u CuTeDSL-u bila je specijalizirana RMS norma jezgra, optimizirana za male dimenzije. QK norma, koja se koristi u sloju samopažnje nedavnih modela Qwen, primjenjuje se po glavi, na 64, 128 ili 256 elemenata. Za razliku od obične RMS norme, koja radi na skrivenim dimenzijama koje premašuju tisuće, zahtijevajući horizontalnu redukciju kroz dijeljenu memoriju i registre, QK norma treba redukciju samo preko warpova, uz mogućnost sažimanja više glava u jedan warp. Ovo je postalo slučaj korištenja u kojem smo se mogli agresivno optimizirati u vrijeme kompilacije kako bismo pronašli najbolje dodjele tokena u warp za sve dimenzije glava i vrste podataka.

RMS norma je prilično jednostavna operacija, iako zahtijeva horizontalnu sumu za izračun kvadrirane norme vektora:

cpp

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

Budući da su glave pažnje obično male, a operacije po niti u velikoj mjeri vektorizirane, jedan warp može izračunati kvadratnu normu. Kako bi se maksimizirala memorijska propusnost, svaka nit treba obraditi 128 bitova ulaza, što se pretvara u 8 bfloat16 vrijednosti. Posljedično, poluwarp je dovoljan za obradu glave, a kvadrirana norma može se odrediti korištenjem horizontalne redukcije izvedene pomoću shuffle_sync_bfly. Polazeći od dimenzije glave i vrste podataka, možemo odrediti broj glava po warpu i specificirati argumente i broj krugova redukcije u vrijeme kompilacije.

Tada možemo izgraditi specijaliziranu jezgru s relativno malo linija koda, oslanjajući se na automatski vektorizirana učitavanja/pohrane i koristeći operacije niske razine specifične za arhitekturu za redukciju warpa. Jezgra određuje na kojoj glavi (Q ili K) treba raditi i identificira odgovarajuće težine koje treba primijeniti. Učitava cijeli vektor iz glave i najprije lokalno zbraja. Nakon toga, na temelju broja niti po glavi, izračunava se kvadratna norma. Na kraju, RMS norma, automatski vektorizirana, se određuje i upisuje natrag u memoriju. Parametri specifični za model, kao što su WEIGHT_BIAS i EPS, sažeti su u konstante.

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 usporedbi s običnom RMS norma jezgrama, ova specijalizirana varijanta pruža 2-3 puta bolju propusnost i latenciju za ispunu i dekodiranje.

MoE slanje i kombiniranje premještaju ulazne tokene stručnjacima kojima su usmjereni, obavljajući ponderiranu akumulaciju nakon što MoE GEMM jezgre izvrše relevantna množenja matrica. Performanse su od posebnog interesa jer se različiti proizvodi oslanjaju na MoE modele različitih veličina. Na temelju vrsta uvođenja razlikujemo različite jezgre koje treba optimizirati:

TP=1: Lokalna jezgra za jedan GPU. Slanje premješta tokene relevantnim stručnjacima, a kombiniranje ih akumulira. Koristi se prvenstveno za uvođenja modela Qwen3 30B na H100.

DP=1, TP ≠ 1: Tenzorski paralelna jezgra s jednim čvorom. Ulaz je repliciran, slanje odabire tokene za trenutni rang, a kombiniranje ih akumulira kroz sve rangove. Akumulacija preko NVLink-a ili NVLS-a. Koristi se za velike uređaje za ispunu.

DP ≠ 1, NVLink: Podatkovno paralelna jezgra s izbornom tenzorskom paralelnošću. Svaki rang šalje tokene svakom drugom rangnom mjestu koje hostira određene stručnjake. Podaci o usmjeravanju se agregiraju, tokeni se prenose preko NVIDIA NVLink-a. Koristi se za podatkovno paralelna uvođenja H200 ili GB200.

DP ≠ 1, InfiniBand: Podatkovno paralelna i tenzorski paralelna jezgra od točke do točke koja radi i preko NVLink-a i preko NVIDIA Quantum InfiniBand-a. Koristi se za uvođenja na više čvorova Hopper.

Lokalna jezgra broji broj tokena dodijeljenih svakom stručnjaku na temelju tablice usmjeravanja i dodjeljuje jedinstveni indeks svakom tokenu unutar popisa tokena koji pripadaju svakom stručnjaku. Nakon koraka usmjeravanja, tokeni se kopiraju u opcionalno popunjene međuspremnike kako bi se proslijedili MoE GEMM jezgri. Koristeći lokacije tokena, jezgra za kombiniranje tada izvodi akumulaciju u registrima, upisujući konačne tokene u izlazni međuspremnik.

Faza slanja razlikuje se između ispune i dekodiranja: za malo tokena ima malo podataka o usmjeravanju za čitanje, pa ih obrađujemo u jednom bloku i signaliziramo ostatak mreže, pri čemu pojedinačni blokovi kopiraju komade tokena. Dodjela tokena određuje se atomskim inkrementiranjem brojača u dijeljenoj memoriji. Budući da ispunjavanje mora obraditi više podataka o usmjeravanju, usmjeravanje je podijeljeno u više blokova u zasebnoj jezgri. Indeksi se prvo lokalno dodjeljuju pomoću brojača u dijeljenoj memoriji, a zatim se agregiraju kroz sve blokove u globalnoj memoriji kako bi se odredio položaj svakog tokena. Mreža ove jezgre ograničena je na broj SM-ova kako bi se mogli koristiti mrežni mehanizmi barijera. Naknadna jezgra obavlja kopiranja, pokrenuta preko znatno većeg broja blokova i warpova kako bi se u potpunosti iskoristila dostupna memorijska propusnost.

Kombinirana jezgra je relativno jednostavna jer izvodi vektoriziranu akumulaciju, oslanjajući se na informacije o usmjeravanju i pomacima koje izračunava slanje. CuTeDSL se ovdje pokazao vrijednim jer se izračun mogao izraziti kroz jednostavne CuTe tenzore koji se kompajliraju u učinkovit kod za sve vrste podataka i skrivene dimenzije.

U tenzorski paralelnoj implementaciji informacije o usmjeravanju i ulazni tokeni su replicirani. Lokalna jezgra za slanje može se u potpunosti ponovno koristiti uz malu promjenu za usmjeravanje i odabir samo tokena dodijeljenih stručnjacima koji se nalaze na trenutnom rangnom mjestu. S nekim parametrima i specijalizacijom u vrijeme kompilacije, lokalne jezgre za slanje se ponovno koriste.

Budući da je broj rangova obično manji od broja stručnjaka kojima je token usmjeren, u prosječnom se slučaju može pretpostaviti da svaki token ima barem jednu od svojih replika usmjerenu na svaki rang. U ovoj situaciji, lokalna kombinirana jezgra može se ponovno koristiti za agregiranje lokalnih kopija tokena, nakon čega slijedi all-reduce operacija zbrajanja kroz rangove koji sudjeluju. U slučaju jezgre za ispunu pokrećemo odvojene jezgre, međutim dekodiranje stapa obje operacije u jedno pokretanje.

Za all-reduce koristimo obični NVLink ako sudjeluju samo dva rangna mjesta, prebacujući se na NVLS kada je uključeno 4 ili više rangnih mjesta radi povećanja propusnosti od 10-20%. Zahvaljujući DSL-u, možemo dijeliti većinu logike redukcije i razmjene vršnjačkih međuspremnika između MoE i redovitih all-reduce jezgri. Ovo je situacija u kojoj je bliska kontrola nad hardverom korisna, budući da su za NVLS potrebne DSL operacije nad multimem.ld_reduce i multimem.st. Sama implementacija redukcije uvelike slijedi primjer all_reduce iz NCCL testova.

Podatkovno paralelne jezgre oslanjaju se na NVLink za agregiranje tokena za više uređaja. U fazi slanja, prije razmjene tokena, parnjaci se najprije lokalno usmjeravaju kako bi identificirali ciljni pomak na svakom rangnom mjestu za svaki token koji šalju. Nakon toga počinju kopirati tokene u međuspremnike za primanje dodijeljene za svakog pošiljatelja kako bi se izbjegli sukobi. Upisi oslobađanja i pohrane po tokenu koriste se za signaliziranje završetka. Na strani primatelja, podaci o usmjeravanju se agregiraju kako bi se svi tokeni mogli sažeti u jedan međuspremnik za MoE GEMM. Primatelji anketiraju zastavicu koju su upisali parnjaci s operacijom nabave i učitavanja kako bi čekali da se upisi dovrše.

Na strani kombiniranja, tokeni se kopiraju preko NVLink-a na izvorne rangove, uz zastavicu po tokenu koja označava primitak. Kombinirana jezgra slična je lokalnoj i jednodevarnoj varijanti, s malim razlikama u adresiranju i nekom dodatnom logikom sinkronizacije.

Izgradili smo knjižnicu PTX omotača za izlaganje učitavanja i pohrana s usvajanjem oslobađanja s opsegom na razini sustava i uređaja za implementaciju ove jezgre. Poboljšavanjem granularnosti sinkronizacijskih zastavica i dodjeljivanjem više blokova za slanje i primanje tokena preko NVLink-a, postigli smo 2x poboljšanja u odnosu na naše prethodne CUDA jezgre.

Naše jezgre od točke do točke prethodno su postigle vrhunske performanse na Hopper GPU-ovima uparenim s ConnectX-7 InfiniBand adapterima. U međuvremenu smo prenijeli kombinirane jezgre na CuTeDSL radi daljnjeg poboljšanja performansi. Budući da su se CUDA verzije već uvelike oslanjale na specijalizaciju u vrijeme kompilacije, pokazale su se kao uska grla vremena kompilacije u našem repozitoriju, prisiljavajući nas da izostavimo neke parametre. Migracijom na DSL kompiliran u trenutku izvršavanja, specijalizirali smo se za više parametara i prilagodili alokacije mreže, uvodeći daljnje razlike između ispune i dekodiranja za otprilike 3% poboljšanja latencije.

Iz perspektive ROSE jezgri, podrška za Blackwell spada u dvije kategorije: jezgre neovisne o arhitekturi i jezgre specifične za arhitekturu. Jezgre specifične za arhitekturu, koje se oslanjaju na tenzorske jezgre, zahtijevaju eksplicitan rad na prilagodbi. Međutim, zahvaljujući korištenju DSL-a, jezgre neovisne o arhitekturi prilagođavaju se bez ikakvih problema.

Blackwell GPU ima do 160 SM-ova, za razliku od Hopperovih 132. Za većinu jezgri ovo je konstanta u vrijeme kompilacije, a logika u kodu pokretanja na strani hosta uzima je u obzir pri određivanju optimalne jezgre za pokretanje (ispuna naspram dekodiranja). Osim toga, Blackwell također uvodi 256-bitne instrukcije učitavanja, koje se lako mogu emitirati povećanjem širine vektora (posljednja dimenzija CuTe tenzora) radi poboljšanja iskoristivosti memorijske propusnosti.

Blackwell uvodi hardversku podršku za značajan broj shema kvantizacije (MXFP8, MXFP4, NVFP4). Podrška za ove sheme implementirana je prvenstveno prilagođavanjem postojećih jezgri za skaliranje po blokovima FP8. Podrška za JIT je ovdje ključna jer se novi načini rada mogu podržati bez eksplozije konfiguracijskog prostora u vrijeme kompilacije. U prošlosti smo otkrili da je Triton bio nespretan u izražavanju horizontalnih redukcija preko nekoliko elemenata, preko polovice ili četvrtine warpa. Fino zrnata kontrola nad redukcijama bila je ključna za postizanje optimalnih performansi pri izračunavanju apsolutne maksimalne vrijednosti preko 16 ili 32 elementa za NVFP4 odnosno MXFP8/MXFP4.

Kao relativno novom jeziku, LLM-ovi su imali malo kontakta s CuTeDSL kodom tijekom svog predtreninga. Dodatno, stil našeg sadašnjeg koda, s jezgrama optimiziranim za aktivacije i norme, bio je malo drugačiji od postojećih primjera CUTLASS ili Quack, koji su bili uvelike optimizirani za TMA i instrukcije množenja matrica specifične za arhitekturu. Dolazak do točke u kojoj LLM-ovi mogu pomoći u pisanju jezgri bio je postupan proces.

Ručno smo izgradili početni korpus jezgri i uslužnih funkcija, prateći dosljedne uzorke, prije nego što smo se mogli osloniti na agente za konverziju jezgri. Kako je korpus napredovao, mogli smo vršiti grube prijevode s CUDA-e na CuTeDSL, uz značajne naknadne ručne preinake. Posebno smo primijetili da su se modeli borili s pretvorbama pokazivača u tenzor i vektorizacijom, zaglavivši u pokušajima ispravljanja pogrešaka tipova kada su ciljali kompajliranje testa. Pokazali su malo razumijevanja mrežne prirode jezgri, brkajući sinkronizaciju niti i mreže. Što je još frustrirajuće, postojala je iskušenje haluciniranja nepostojećih prijevoda knjižnica specifičnih za CUDA-u umjesto izgradnje na našim vlastitim zamjenama za CuTeDSL. S vremenom, kako je korpus CuTeDSL jezgri rastao, otkrili smo da LLM-ovi mogu izgraditi većinu kosturnog koda i popuniti praznine, dopuštajući nam da se usredotočimo na dizajn visoke razine.

U usporedbi s CUDA-om, otkrivamo da CuTeDSL pruža puno bolje performanse u vrijeme kompilacije, omogućujući nam specijalizaciju za više parametara i postizanje poboljšanja performansi. Iako su jezgre nešto opširnije zbog nedostatka određenih apstrakcija u usporedbi s Tritonom, smatramo da je složenost opravdana jer nam fino zrnata kontrola nad hardverom omogućuje izgradnju bržih jezgri.

U vrijeme pisanja ovog teksta, jaz u performansama između CUTLASS C++ i CuTeDSL jezgri minimalan je, otvarajući uzbudljive prilike za sastavljanje i stapanje visokoučinkovitih GEMM jezgri s komunikacijom između uređaja, normalizacijom ili aktivacijama. Obvezivanjem na ovaj DSL moći ćemo izvući dodatne performanse iz naših uređaja kako bismo poboljšali propusnost kroz naše radne zadatke, od ugrađivanja 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}
}

Zaključivanje u tvrtki Perplexity AI

Zašto CuTeDSL?

Iskustvo razvojnih inženjera

Veliki konfiguracijski prostor

Unaprijed ispunite i dekodirajte specijalizaciju

Mrežna sinkronizacija

CuTeDSL jezgre

QK norma

Miješanje stručnjaka (MoE) - slanje/kombiniranje

Lokalno

Jedan čvor

Podatkovni paralelizam

Od točke do točke (Peer-to-Peer)

Podrška za Blackwell

Generiranje jezgri potpomognuto umjetnom inteligencijom

Zaključak i budući rad

Reference

Citiranje