CuTeDSL la Perplexity

La Perplexity, folosim motorul intern Runtime-Optimized Serving Engine (ROSE) pentru a rula modele, de la încorporări (embeddings) până la modele lingvistice de mari dimensiuni (LLM-uri) cu un trilion de parametri. Situat în spatele tuturor API-urilor, inclusiv Sonar, Search și Embeddings, motorul nostru de execuție a fost creat pentru a le oferi cercetătorilor posibilitatea de a explora cele mai potrivite modele și arhitecturi, implementând totodată rapid cei mai buni candidați în producție. Pentru a realiza acest lucru, trebuie să fim capabili să implementăm, să lansăm și să optimizăm inferența pentru arhitecturi de ultimă generație într-un ritm rapid. În această postare pe blog, explorăm modul în care nucleele CuTeDSL integrate în motorul nostru de inferență ne permit să construim nucleele GPU necesare pentru a aduce modelele la performanță maximă pe GPU-urile NVIDIA.

Modelele personalizate, construite de diferite echipe de la Perplexity, sunt găzduite intern pe GPU-uri NVIDIA Hopper și Blackwell folosind motorul de inferență ROSE, construit de echipa de inferență AI. ROSE a fost creat inițial pentru a deservi mai bine modele Llama personalizate atât pentru decodarea limbajului, cât și pentru clasificare, având o interfață compatibilă cu NVIDIA Triton Inference Server. În timp, a evoluat într-un motor de inferență complet care deservește LLM-uri substanțiale, precum și modele de clasificare, scorare și alte modele bazate pe transformere pentru o gamă largă de echipe, alimentând produsele Perplexity. Scopul de proiectare al ROSE este de a fi simplu: în timp, se investește efort pentru a atinge performanța maximă, însă prioritatea principală este adaptarea la arhitecturile și cazurile de utilizare în continuă evoluție pentru modelele noastre.

Spațiul de configurare al kernelului pentru tipuri de date, dimensiuni ale capetelor, constante de compilare și setări de lansare.

La un nivel înalt, ROSE constă dintr-un motor care adaptează un model la o interfață cu care clienții pot interacționa. Motoarele sunt extrem de reutilizabile, gestionând programarea și gruparea (batching) cererilor și preluând tokenurile decodate, încorporările sau scorurile de la modelele subiacente. Inițializarea dispozitivului, canalele de comunicare între procese și noduri, precum și încărcarea ponderilor sunt toate gestionate de acestea. Pentru LLM-uri, motorul este responsabil pentru grupare (batching), fragmentare (chunking) și eșantionare. De asemenea, este responsabil pentru alocarea stocării KV, atât pentru atenția completă, cât și pentru cea liniară, precum și pentru potrivirea prefixelor. Motoarele de încorporare sunt mai simple, deoarece efectuează doar grupare online. Execuția este optimizată astfel încât motorul pregătește runda următoare de lucru în timp ce modelul rulează în paralel pe accelerator.

ROSE expune o colecție bogată de straturi personalizate pentru modelare, care împachetează nuclee GPU construite cu CuTeDSL, Triton, CUDA, CUTLASS și cuBLAS. Pentru multiplicarea de matrice și implementările de atenție, ne bazăm în principal pe nuclee NVIDIA din CUTLASS și cuBLAS. Restul conductei, inclusiv nucleele de încorporare, normă, MoE și activare, fac parte din biblioteca de nuclee a ROSE. Pentru a susține optim hardware-ul pe care rulăm inferența, ROSE este echipat cu propriile primitive de comunicare între dispozitive pentru rutarea, dispecerizarea și combinarea MoE. Deși celelalte nuclee sunt relativ simple, ele trebuie specializate temeinic pentru a obține performanțe maxime pe toate dimensiunile ascunse diferite și pe dispozitivele pe care rulează. În plus, ne bazăm pe specializarea la momentul compilării pentru a acomoda toate opțiunile și ajustările diferitelor modele (bias pentru ponderi față de lipsa biasului pentru norma RMS, funcții de activare etc.). În consecință, am decis să folosin DSL-uri pentru a le implementa: inițial am construit ROSE folosind Triton. Cu toate acestea, de atunci am adoptat CuTeDSL ca mediu principal pentru programarea GPU.

Din perspectivă matematică, nucleele de inferență sunt de obicei nesofisticate. În consecință, la dezvoltarea unui nucleu, complexitatea nu constă în general în exprimarea calculului de efectuat, ci în modul exact în care hardware-ul ar trebui să îl execute. Pentru a construi cele mai bune nuclee, avem nevoie de un limbaj care să ofere abstracțiunile corecte asupra hardware-ului subiacent. Limbajul corect ne permite să specificăm operațiunile de nivel înalt pentru a efectua calculul, specializându-le și optimizându-le agresiv în toate punctele dintr-un spațiu de configurare larg.

CuTeDSL este un DSL bazat pe Python pentru GPU-uri NVIDIA, construit pe baza algebrei de machetare CuTe și MLIR. Acesta se compilează just-in-time în cod PTX înalt optimizat, menținând în același timp controlul asupra primitivelor hardware de nivel inferior. Scopul său este de a oferi aceleași caracteristici și performanțe ca și CUTLASS C++, fără timpii de compilare costipizi care împiedică viteza de dezvoltare. Dintre limbajele pe care le-am analizat sau cu care am lucrat (inclusiv CUDA, cuTile, TVM IR și Triton), îl considerăm cel mai potrivit pentru dezvoltarea nuclerelor noastre.

În general, deși nu este la fel de nivel înalt și concis ca Triton sau cuTile, am constatat că algebra de machetare CuTe Layout Algebra și compozițiile sale (expuse prin cute.Tensor) au fost substituenți rezonabili. În plus, CuTeDSL oferă acces nerestricționat la primitivele hardware, oferindu-ne nivelul de control de care avem nevoie pentru a obține performanță de vârf în toate cazurile limită (edge cases). Cu CUDA, oportunitățile de specializare la momentul compilării au fost limitate, deoarece fiecare parametru suplimentar a adăugat un cost nerezonabil timpii de compilare sau complexitate sistemului de construcție. În schimb, cu un DSL compilat Just-In-Time, putem specializa mai agresiv fără a sacrifica ritmul de iterare. Spre deosebire de TVM IR, care folosea CUDA ca pas intermediar, am considerat că conducta bazată pe MLIR care coboară la PTX este mai robustă și mai ușor de depanat decât una care emitea CUDA ca pas intermediar. În plus, prin utilizarea de op-uri personalizate pentru a împacheta instrucțiunile PTX, am constatat că putem valorifica pe deplin toate caracteristicile hardware-ului subiacent.

Cea mai mare parte a costului de compilare pentru nucleele CUDA este de obicei cheltuită în faza de expansiune a șabloanelor (template expansion) a nuclerelor, deoarece optimizarea și compilarea din stadiul PTX încolo pot fi finalizate în câteva secunde. Cu mulți parametri de compilare, această etapă a compilării poate exploda la câteva minute pentru un nucleu non-trivial. Acest blocaj este accentuat de faptul că oportunitățile de compilare paralelă apar doar mai târziu, la nivelurile IR și PTX, gestionarea șabloanelor C++ având loc secvențial. Deși există trucuri pentru a rezolva acest lucru, prin mutarea instanțierilor de șabloane în unități de compilare separate, acestea necesită ajustări semnificative ale sistemelor de construcție și utilizarea unor macrocomenzi greoi. În schimb, nucleele CuTeDSL care sunt compilate Just-In-Time (JIT) nu necesită reguli de construcție complexe și nu îngreunează timpii de compilare, deoarece ciclul de dezvoltare poate fi axat pe un singur punct din spațiul de configurare, în timp ce testele unitare atente ulterioare pot acoperi întregul spațiu de proiectare al nucleului.

Chiar dacă compilarea JIT mută costurile de la faza de dezvoltare la timpul de execuție, suprasarcina este mai mult decât acceptabilă. Atunci când un model este implementat, acesta activează doar câteva dintre punctele din întregul spațiu de configurare, generând suprasarcini de ordinul secundelor în timpul pornirii/încălzirii motorului de inferență. În schimb, pre-compilarea tuturor configurațiilor unui nucleu ar putea dura câteva minute. În plus, suprasarcinele pot fi atenuate și mai mult prin stocarea în cache a artefactelor compilate prin JIT.

Costul mare de compilare al CUDA este accentuat de dificultatea de a urmări originea erorilor atunci când se utilizează dispecerizarea și specializarea bazate pe șabloane (template-based), așa cum este cazul cu CuTe și CUTLASS. Cu un DSL, erorile pot fi indicate către linia de cod care le-a generat, ușurând semnificativ depanarea. În plus, este mult mai simplu să se mapeze codul generat la instrucțiunea DSL care l-a generat prin intermediul stadelor MLIR intermediare sau bazându-se pe informațiile de depanare atașate la PTX-ul generat. Datorită acestor aspecte ale DSL-ului, am constatat că suntem mult mai productivi în dezvoltarea nuclerelor.

Deși nucleele au multe similatități între modele, ele trebuie totuși echipate cu o serie de butoane de configurare care pot avea un impact asupra performanței. Parametrii relevanți includ:

Dimensiune ascunsă / Dimensiune cap / Număr de capete de atenție (Hidden Dimension/Head Dimension/Number of Attention Heads)

Tipuri de date de intrare/ieșire

Schemă de cuantizare

Constante de compilare (biasuri, funcții de activare etc.)

Arhitectura ROSE, inclusiv motorul său, modelul, biblioteca kernelului și serviciile.

Specializarea la momentul compilării este utilă atât pentru performanță, cât și pentru depanare. În special într-un procesor cu conducte (pipelined), cum este cazul GPU-urilor, o ramură inexistentă este mai bună decât o ramură luată mereu sau niciodată luată. În plus, compilarea unui nucleu pentru un set specific de opțiuni de configurare are ca rezultat mai puține instrucțiuni PTX/SASS, care sunt mai ușor de analizat și depanat atunci când se utilizează cuda-gdb în timpul unei investigații.

Pentru o dimensiune de lot (batch size) și o configurație date, se pot estima traficul așteptat către DRAM și numărul de operații aritmetice. Ulterior, se pot lua decizii optime de la caz la caz pentru alocarea optimă a grilei, sincronizare și utilizarea anumitor caracteristici hardware. CuTeDSL ne permite să definim și să reutilizăm componente pentru o specializare profundă fără a genera un cost de compilare inacceptabil.

În general, performanța nuclerelor de inferență trebuie optimizată pentru două cazuri de utilizare distincte:

Prefill: dimensiune de lot mică (puține fragmente, secvențe), multe tokenuri per secvență. Debitul (throughput) trebuie maximizat atât pentru încorporări, cât și pentru tokenurile de intrare ale LLM-urilor. Latență de la zeci la sute de milisecunde.

Decode: dimensiuni de lot de la mici la mari, de la unu (fără decodare speculativă) la câteva tokenuri per secvență (speculație țintă-schiță, predicție cu mai multe tokenuri). Latență apropiată de latența de lansare a nucleului.

Pentru GEMM și atenție, etapa de preumplere (prefill) rulează pe un număr mare de tokenuri, iar performanța este limitată de debitul de calcul al dispozitivului. Decodarea implică aceeași operațiune pe dimensiuni de lot mici, timpii de execuție fiind dominați de costul încărcării activărilor și a ponderilor din DRAM în multiprocesoarele de flux și nucleele tensoriale, fiind constrânși în schimb de lățimea de bandă a memoriei dispozitivului. Chiar dacă dimensiunea lotului sau numărul total de tokenuri dintr-un lot este un indicator relativ clar, pragul de comutare între prefill și decode este mai subtil, deoarece lățimea de bandă a memoriei este determinată și de dimensiunea ascunsă a modelului și de tipul de date al activării. În consecință, nucleele sunt împărțite în implementări de prefill și decode dacă diferențele algoritmice sau atribuirile de grilă/bloc favorizează vreunul dintre ele.

Parametrul care diferă universal între nucleele de prefill și decode este numărul de warp-uri dintr-un bloc (numărul de fire de execuție împărțit la dimensiunea warp-ului de 32 de fire de execuție per warp). Este de dorit ca acest parametru să fie o constantă la momentul compilării, deoarece poate ajuta la alocarea registrelor. În CUDA, acest lucru se face de obicei prin specificarea numărului de fire de execuție prin __launch_bounds__ și furnizarea numărului de warp-uri ca parametru de șablon:

cpp

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

}

Numărul de registre disponibile pentru toate firele de execuție dintr-un bloc este o constrângere hardware, de 64K pe Hopper și Blackwell. Dacă numărul de fire de execuție este cunoscut din timp, compilatorul știe exact câte registre poate aloca fiecărui fir de execuție și poate reuși să compileze nuclee optimizate mai agresiv. Cu toate acestea, acest lucru vine cu un cost în timpul compilării, deoarece acum nucleul trebuie re-compilat pentru diferite dimensiuni de warp. Cu CuTeDSL, aceasta nu mai este o problemă, deoarece nu trebuie să plătim costul de compilare pentru întregul spațiu de configurare în avans. De asemenea, DSL-ul se așteaptă ca dimensiunea blocului să fie o expresie constantă, efectuând implicit optimizarea implicată de __latch_bounds__ în CUDA.

Pe lângă transformarea numărului de warp-uri într-o constantă, este de dorit să se varieze acesta între prefill și decode. Pentru a utiliza pe deplin lățimea de bandă a memoriei unui GPU, nu este suficient să alocați întotdeauna un număr mare de warp-uri pe bloc, deoarece lățimea de bandă maximă poate fi obținută doar prin utilizarea a cât mai multor Multiprocesoare de Flux (SM) disponibile. În consecință, pentru un număr dat de tokenuri și tipuri de date și dimensiuni ascunse cunoscute, alegem mai întâi o dimensiune vectorială ideală per fir de execuție (8 pentru bfloat16, 4 pentru float32) pentru a ne asigura că codul generat utilizează instrucțiunile de citire/scriere (load/store) pe 128 de biți (256 de biți pe Blackwell). Cunoscând dimensiunea vectorului, putem calcula numărul de warp-uri cerut de nucleu. Încercăm mai întâi să maximizăm numărul de SM-uri: dacă există mai puține warp-uri decât SM-uri, este ales un singur warp per bloc. Aceasta este de obicei varianta de decodare (decode) a nucleului. Dacă există suficiente warp-uri, mărim treptat numărul de warp-uri la 8, 16 sau 32, evaluând o limită superioară după care de obicei nu mai vedem o îmbunătățire a performanței.

Deși Triton a efectuat unele dintre aceste optimizări pe sub capotă, am constatat că putem obține rezultate similare utilizând cu atenție algebra de machetare CuTe, fără a fi constrânși de neajunsurile unui compilator de nivel înalt. În special atunci când sunt implicate reduceri orizontale sub-warp, am constatat că CuTeDSL ne-a oferit controlul fin necesar pentru a genera cod optim.

În unele situații, pentru a distribua lucrul pe cât mai multe blocuri posibil, cooperative_groups::this_grid().sync() poate fi utilizat pentru a sincroniza toate blocurile participante. Se aplică restricții: deoarece un bloc nu poate fi anulat de la programare dintr-un SM odată ce începe execuția, pentru a garanta progresul înainte, numărul de blocuri nu poate depăși numărul de SM-uri disponibile pe dispozitiv, asigurându-se că în orice moment toate blocurile pot actualiza și interoga locațiile de memorie globală implicate în sincronizare. Această restricție este impusă de cudaLaunchCooperativeKernel. Costul unei bariere de grilă este de aproximativ 3us pe Hopper, plus costul golirii tuturor operațiunilor de memorie anterioare printr-o barieră de memorie implicită în anumite situații. Per total, acest lucru este comparabil cu suprasarcina împărțirii unui nucleu în două lansări, ducând la o decizie de proiectare sensibilă la performanță. În ciuda faptului că Programmatic Dependent Launch (PDL) poate reduce suprasarcina de lansare începând cu Hopper, pentru nucleele de decodare, evitarea suprasarcinii merită.

Dacă una dintre sarcinile de lucru necesită mai multe blocuri decât numărul de SM-uri, nucleul trebuie împărțit, bazându-se pe bariera implicită a lansărilor de nuclee. Acesta este în principal cazul pentru prefill. Dacă numărul de blocuri care este ideal pentru toate sarcinile de lucru se încadrează în grilă, este preferată o barieră de grilă, așa cum este cazul nuclerelor de decodare care operează pe un număr mai mic de tokenuri.

În CuTeDSL, aceste decizii pot fi luate în codul de lansare din partea gazdei (host-side), iar nucleele pot fi compuse și intercalate cu bariere de grilă după cum este necesar, așa cum se ilustrează în exemplul de mai jos. Cu o astfel de strategie hibridă, nucleele de decodare pot fi optimizate pentru a atinge latențe minime, iar cele de prefill pot fi optimizate pentru un debit maximal. Spre deosebire de soluțiile bazate pe CUDA, această logică poate fi definită mai concis și cu o suprasarcina mai mică a timpului de compilare.

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

În cazul nuclerelor noastre MoE all-to-all, așa cum este detaliat mai jos, flexibilitatea de a construi nuclee cu sau fără bariere de grilă și de a le optimiza la momentul compilării ne-a ajutat să obținem câștiguri substanțiale de performanță prin CuTeDSL.

Primul nucleu pe care l-am construit în CuTeDSL a fost un nucleu de normă RMS specializat, optimizat pentru dimensiuni mici. Norma QK, utilizată în stratul de auto-atenție al modelelor Qwen recente, este aplicată pe bază de cap (per-head), pe 64, 128 sau 256 elemente. Spre deosebire de norma RMS obișnuită, care operează pe dimensiuni ascunse care depășesc mii, necesitând o reducere orizontală atât prin memoria partajată, cât și prin registre, norma QK necesită o reducere numai între warp-uri, cu posibilitatea de a înghesui mai multe capete într-un singur warp. Acesta a devenit un caz de utilizare în care am putut optimiza agresiv la momentul compilării pentru a găsi cele mai bune atribuiri de la token la warp pentru toate dimensiunile de cap și tipurile de date.

Norma RMS este o operațiune destul de simplă, deși necesită o sumă orizontală pentru a calcula norma pătrată a unui vector:

cpp

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

Deoarece capetele de atenție sunt de obicei mici și operațiunile per fir de execuție sunt puternic vectorizate, un singur warp poate calcula norma pătrată. Pentru a maximiza lățimea de bandă a memoriei, fiecare fir de execuție ar trebui să proceseze 128 de biți de intrare, ceea ce se traduce prin 8 valori bfloat16. În consecință, un jumătate de warp (half-warp) este suficient pentru a procesa un cap, iar norma pătrată poate fi determinată folosind o reducere orizontală efectuată de shuffle_sync_bfly. Pornind de la dimensiunea capului și tipul de date, putem determina numărul de capete per warp și putem specifica argumentele și numărul de runde de reducere în momentul compilării.

Putem apoi construi un nucleu specializat cu relativ puține linii de cod, bazându-ne pe încărcări/scrieri vectorizate automat și folosind operațiuni specifice arhitecturii de nivel inferior pentru reducerea warp-ului. Nucleul determină pe ce cap (Q sau K) să lucreze și identifică ponderile corespunzătoare de aplicat. Încarcă un vector întreg din cap și mai întâi însumează la nivel local. Ulterior, în funcție de numărul de fire de execuție per cap, se calculează norma pătrată. În cele din urmă, norma RMS, vectorizată automat, este determinată și scrisă înapoi în memorie. Parametrii specifici modelului, cum ar fi WEIGHT_BIAS și EPS, sunt pliați pe constante (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)

În comparație cu un nucleu de normă RMS obișnuit, această variantă specializată oferă un debit și o latență de 2-3 ori mai bune atât pentru prefill, cât și pentru decode.

Dispecerizarea și combinarea MoE amestecă tokenurile de intrare către experții la care sunt rutate, efectuând o acumulare ponderată după ce nucleele MoE GEMM efectuează multiplicările de matrice relevante. Performanța este de un interes deosebit, deoarece diverse produse se bazează pe modele MoE de diferite dimensiuni. În funcție de tipurile de implementare, distingem diferite nuclee de optimizat:

TP=1: Nucleu local pentru un singur GPU. Dispecerizarea amestecă tokenurile către experții relevanți, iar combinarea le acumulează. Utilizat în principal pentru implementările Qwen3 30B pe H100.

DP=1, TP ≠ 1: Nucleu tensor-paralel, cu un singur nod (singlenode). Intrarea este replicată, dispecerizarea selectează tokenurile pentru rangul curent, iar combinarea le acumulează pe toate raturile (ranks). Acumulare prin NVLink sau NVLS. Utilizat pentru prefill-ere mari.

DP ≠ 1, NVLink: Nucleu data-paralel cu paralelism de tensori opțional. Fiecare rang trimite tokenuri către fiecare alt rang care găzduiește experți specifici. Informațiile de rutare sunt agregate, tokenurile sunt transferate prin NVIDIA NVLink. Utilizat pentru implementări data-paralele H200 sau GB200.

DP ≠ 1, InfiniBand: Nucleu Peer-to-Peer data-paralel și tensor-paralel care operează atât prin NVLink, cât și prin NVIDIA Quantum InfiniBand. Utilizat pentru implementări Hopper cu mai multe noduri.

Nucleul local numără numărul de tokenuri atribuite fiecărui expert pe baza tabelului de rutare și atribuie un index unic fiecărui token din lista de tokenuri aparținând fiecărui expert. După pasul de rutare, tokenurile sunt copiate în tampoane (buffers) opțional umplute (padded) pentru a fi transmise mai departe către un nucleu MoE GEMM. Folosind locațiile tokenurilor, nucleul de combinare efectuează apoi o acumulare în registre, scriind tokenurile finale într-un tampon de ieșire.

Faza de dispecerizare variază între prefill și decode: pentru puține tokenuri, există puține informații de rutare de citit, așa că le procesăm într-un singur bloc și semnalizăm restul grilei, cu blocuri individuale care copiază fragmente de tokenuri. Atribuirea tokenurilor este determinată prin incrementarea atomică a contoarelor din memoria partajată. Deoarece prefill trebuie să proceseze mai multe informații de rutare, rutarea este împărțită pe mai multe blocuri într-un nucleu separat. Indicii sunt mai întâi atribuiți local folosind contoare în memoria partajată și sunt apoi agresați (aggregated) pe toate blocurile din memoria globală pentru a determina poziția fiecărui token. Grila acestui nucleu este limitată (clamped) la numărul de SM-uri pentru a putea utiliza bariere la nivel de grilă. Un nucleu subsecvent efectuează copiile, lansat pe un număr mult mai mare de blocuri și warp-uri pentru a utiliza pe deplin lățimea de bandă a memoriei disponibile.

Nucleul de combinare este relativ simplu, deoarece efectuează acumulare vectorizată, bazându-se pe informațiile de rutare și pe offseturile calculate prin dispecerizare. CuTeDSL s-a dovedit valoros aici, deoarece calculul putea fi exprimat prin tensori CuTe simpli care se compilează în cod eficient pentru toate tipurile de date și dimensiunile ascunse.

Într-o implementare tensor-paralelă, informațiile de rutare și tokenurile de intrare sunt replicate. Nucleul de dispecerizare local poate fi complet reutilizat cu o ușoară modificare pentru a ruta și selecta numai tokenurile atribuite experților aflați pe rangul curent. Cu unii parametri și specializare la momentul compilării, nucleele de dispecerizare locală sunt reutilizate.

Deoarece numărul de raturi (ranks) este de obicei mai mic decât numărul de experți la care este rutat un token, în caz mediu se poate presupune că fiecare token are cel puțin una dintre replici rutată către fiecare rang. În această situație, nucleul de combinare local poate fi reutilizat pentru a agrega copiile locale ale unui token, urmat de un all-reduce pentru a însuma pe raturile participante. În cazul nucleului de prefill, lansăm nuclee separate, cu toate acestea, decode fuzionează ambele operațiuni într-o singură lansare.

Pentru all-reduce, folosim NVLink simplu dacă participă doar două raturi, trecând la NVLS când sunt implicate 4 sau mai multe raturi pentru o creștere a debitului cu 10-20%. Datorită DSL-ului, putem partaja cea mai mare parte a logicii de reducere și schimburile de tampoane peer între nucleul MoE și nucleele all-reduce obișnuite. Aceasta este o situație în care controlul strâns asupra hardware-ului este benefic, deoarece operațiunile DSL pe multimem.ld_reduce și multimem.st sunt necesare pentru NVLS. Implementarea reducerii în sine urmează în mare măsură exemplul all_reduce din testele NCCL.

Nucleele Data Parallel se bazează pe NVLink pentru a agrega tokenuri pentru mai multe dispozitive. În faza de dispecerizare, înainte de schimbul de tokenuri, perechile (peers) se rutează mai întâi la nivel local pentru a identifica offsetul țintă pe fiecare rang pentru fiecare token pe care îl trimit. Ulterior, încep să copieze tokenurile în tampoane de recepție alocate pentru fiecare expeditor pentru a evita conflictele. Scrierile release-store per token sunt folosite pentru a semnala finalizarea. La capătul receptor, informațiile de rutare sunt agregate pentru a putea condensa toate tokenurile într-un singur tampon pentru MoE GEMM. Receptorii interoghează (poll) fanionul scris de pereche cu un acquire-load pentru a aștepta finalizarea scrierilor.

În partea de combinare, tokenurile sunt copiate prin NVLink către raturile sursă, cu un fanion per token care indică primirea. Nucleul de combinare este similar cu variantele locale și pentru un singur nod, cu mici diferențe la adresare și unele logici de sincronizare suplimentare.

Am construit o bibliotecă de wrapper-e PTX pentru a expune citiri și scrieri release-acquire cu domeniu (scoped) atât la nivel de sistem, cât și de dispozitiv pentru a implementa acest nucleu. Îmbunătățind granularitatea fanioanelor de sincronizare și alocând mai multe blocuri pentru a trimite și a primi tokenuri prin NVLink, am obținut îmbunătățiri de 2x față de nucleele noastre CUDA anterioare.

Nucleele noastre Peer-to-Peer au obținut anterior performanțe de ultimă generație pe GPU-uri Hopper asociate cu adaptoare InfiniBand ConnectX-7. De atunci, am portat nucleele de combinare pe CuTeDSL pentru câștiguri suplimentare de performanță. Deoarece versiunile CUDA se bazau deja puternic pe specializarea la momentul compilării, ele s-au dovedit a fi blocajele în timpul compilării din repository-ul nostru, obligându-ne să omitem unii parametri. Migrând la DSL-ul compilat JIT, ne-am specializat pentru mai mulți parametri și am ajustat alocările de grilă, introducând diferențe suplimentare între prefill și decode pentru o creștere de aproximativ 3% a latenței.

Din perspectiva nuclerelor ROSE, suportul Blackwell se încadrează în două categorii: nuclee agnostice la arhitectură și nuclee specifice arhitecturii. Nucleele specifice arhitecturii, care se bazează pe nuclee Tensor, necesită lucru explicit de adaptare. Cu toate acestea, datorită utilizării DSL-ului, nucleele agnostice la arhitectură se adaptează în mod trivial.

Un GPU Blackwell are până la 160 de SM-uri, spre deosebire de cele 132 ale lui Hopper. Pentru majoritatea nuclerelor, aceasta este o constantă la momentul compilării, iar logica din codul de lansare din partea gazdei ia acest lucru în considerare pentru determinarea nucleului optim (prefill vs decode) de rulat. În plus, Blackwell introduce și instrucțiuni de încărcare pe 256 de biți, care pot fi emise cu ușurință prin mărirea lățimii vectorului (ultima dimensiune a tensorilor CuTe) pentru a îmbunătăți utilizarea lățimii de bandă a memoriei.

Blackwell introduce suport hardware pentru un număr semnificativ de scheme de cuantizare (MXFP8, MXFP4, NVFP4). Suportul pentru acestea a fost implementat în principal prin ajustarea nuclerelor existente pentru scalarea pe blocuri FP8. Suportul JIT este crucial aici, deoarece noile moduri pot fi acceptate fără a exploda spațiul de configurare de la momentul compilării. În trecut, am considerat că Triton este stângaci în exprimarea reducerilor orizontale pe câteva elemente, pe jumătate sau un sfert dintr-un warp. Un control mai fin asupra reducerilor a fost crucial pentru a obține performanțe optime la calcularea unei valori maxime absolute pe 16 sau 32 de elemente pentru NVFP4 și, respectiv, MXFP8/MXFP4.

Fiind un limbaj relativ nou, LLM-urile au avut o expunere redusă la codul CuTeDSL în timpul pre-antrenamentului lor. În plus, stilul codului nostru actual, cu nuclee optimizate pentru activări și norme, a fost ușor diferit față de exemplele CUTLASS sau Quack existente, care au fost puternic optimizate pentru instrucțiuni TMA și de multiplicare de matrice specifice arhitecturii. Atingerea unui punct în care LLM-urile ar putea ajuta la scrierea nucleelor a fost un proces treptat.

Am construit manual un corpus inițial de nuclee și funcții de utilitate, urmând modele consistente, înainte de a ne putea baza pe agenți pentru conversia nuclerelor. Pe măsură ce corpusul a progresat, am putut efectua traduceri brute de la CUDA la CuTeDSL, urmate de o muncă manuală semnificativă. În particular, am constatat că modelele s-au luptat cu conversiile de la pointer la tensor și cu vectorizarea, blocându-se încercând să repare erorile de tip atunci când urmăreau să compileze un test. Au arătat o înțelegere redusă a naturii bazate pe grilă a nuclerelor, confundând sincronizarea firelor de execuție cu cea a grilei. Cel mai frustrant a fost tentația de a halucina traduceri inexistente ale bibliotecilor specifice CUDA în loc să construiască pe baza propriilor noastre înlocuitori CuTeDSL. În timp, pe măsură ce corpusul de nuclee CuTeDSL a crescut, am constatat că LLM-urile pot construi cea mai mare parte a codului schelet și pot completa spațiile libere, permițându-ne să ne concentrăm pe designul de nivel înalt.

În comparație cu CUDA, considerăm că CuTeDSL oferă performanțe mult mai bune la momentul compilării, permițându-ne să ne specializăm pentru mai mulți parametri și să obținem îmbunătățiri de performanță. Deși nucleele sunt puțin mai verboase din cauza lipsei anumitor abstracțiuni în comparație cu Triton, considerăm că complexitatea este justificată, deoarece controlul mai fin asupra hardware-ului ne permite să construim nuclee mai rapide.

La momentul redactării acestui articol, decalajul de performanță dintre CUTLASS C++ și nucleele CuTeDSL este minim, deschizând oportunități incitante de a compune și fuziona nuclee GEMM de înaltă performanță cu comunicare între dispozitive, normalizare sau activări. Angajându-ne față de acest DSL, vom putea extrage performanțe suplimentare de la dispozitivele noastre pentru a îmbunătăți debitul pentru sarcinile noastre de lucru, de la încorporări la LLM-uri.

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

Inferența la Perplexity AI

De ce CuTeDSL?

Experiența dezvoltatorului

Spațiu de configurare mare

Precompletarea și decodificarea specializării

Sincronizarea grilei (Grid synchronization)

Nuclee CuTeDSL

Normă QK (QK Norm)

Dispecerizare/Combinare Amestec de Experți (Mixture-of-Experts Dispatch/Combine)

Local

Un singur nod (Singlenode)

Data Parallel (Paralel pe date)

Peer-to-Peer

Suport Blackwell

Generare de nuclee asistată de AI (AI-Assisted Kernel Generation)

Concluzie și lucrări viitoare

Referințe

Citare