CuTeDSL at Perplexity

V spoločnosti Perplexity používame vlastný engine Runtime-Optimized Serving Engine (ROSE) na obsluhu modelov od vkladov (embeddings) až po LLM s biliónmi parametrov. Náš runtime, ktorý stojí za všetkými rozhraniami API vrátane Sonar, Search a Embeddings, bol vytvorený s cieľom umožniť výskumníkom preskúmať najvhodnejšie modely a architektúry a zároveň rýchlo nasadiť najlepších kandidátov do produkcie. Aby sme to dosiahli, musíme byť schopní rýchlo implementovať, nasadzovať a optimalizovať inferenciu pre najmodernejšie architektúry. V tomto blogovom príspevku sa venujeme tomu, ako jadrá CuTeDSL integrované do nášho inferenčného enginu umožňujú vytvárať potrebné jadrá GPU na dosiahnutie špičkového výkonu modelov na GPU NVIDIA.

Vlastné modely, ktoré vytvorili rôzne tímy v Perplexity, sú interne hostované na GPU NVIDIA Hopper a Blackwell pomocou inferenčného enginu ROSE vytvoreného tímom pre AI inferenciu. Engine ROSE bol pôvodne vytvorený na lepšiu obsluhu prispôsobených modelov Llama na dekódovanie jazyka aj klasifikáciu s rozhraním kompatibilným s NVIDIA Triton Inference Server. Postupne sa vyvinul v plnohodnotný inferenčný engine, ktorý obsluhuje rozsiahle LLM, ako aj modely na hodnotenie, klasifikáciu, skórovanie a iné modely založené na transformeroch pre širokú škálu tímov, ktoré poháňajú produkty spoločnosti Perplexity. Cieľom návrhu ROSE je jednoduchosť: postupne sa investuje úsilie do dosiahnutia špičkového výkonu, pričom hlavnou prioritou je prispôsobiť sa neustále sa vyvíjajúcim architektúram a prípadom použitia pre naše modely.

Priestor konfigurácie jadra naprieč dátovými typmi, rozmermi hláv, konštantami prekladania a nastaveniami spustenia.

Na vysokej úrovni sa ROSE skladá z enginu, ktorý prispôsobuje model rozhraniu, s ktorým môžu klienti komunikovať. Engine je vysoko opakovane použiteľný, riadi plánovanie a dávkovanie požiadaviek a získava dekódované tokeny, vklady (embeddings) alebo skóre zo základných modelov. Zabezpečuje inicializáciu zariadení, komunikačných kanálov medzi procesmi a uzlami a načítavanie váh. Pre LLM je engine zodpovedný za dávkovanie, delenie na bloky (chunking) a vzorkovanie. Zodpovedá tiež za alokáciu úložiska KV pre úplnú aj lineárnu pozornosť, ako aj za zhody predpôn (prefix matching). Motory pre vklady (embedding engines) sú jednoduchšie, pretože vykonávajú iba online dávkovanie. Vykonávanie je optimalizované tak, že engine pripravuje ďalšie kolo práce, zatiaľ čo model beží paralelne na akcelerátore.

ROSE sprístupňuje bohatú kolekciu vlastných vrstiev na modelovanie, ktoré obaľujú jadrá GPU vytvorené pomocou CuTeDSL, Triton, CUDA, CUTLASS a cuBLAS. Pri implementáciách násobenia matíc a pozornosti sa spoliehame predovšetkým na jadrá NVIDIA z CUTLASS a cuBLAS. Zvyšok pipeline vrátane jadier pre vklady (embedding), normalizáciu (norm), MoE a aktivačné jadrá je súčasťou knižnice jadier ROSE. Aby sme čo najlepšie podporili hardvér, na ktorom spúšťame inferenciu, ROSE je vybavený vlastnými primitívami komunikácie medzi zariadeniami na smerovanie, odosielanie a kombinovanie MoE. Zatiaľ čo ostatné jadrá sú relatívne jednoduché, musia byť dôkladne špecializované, aby sa dosiahol špičkový výkon vo všetkých rôznych skrytých dimenziách a zariadeniach, na ktorých bežia. Okrem toho sa spoliehame na špecializáciu v čase kompilácie, aby sme prispôsobili všetky ovládacie prvky a úpravy rôznych modelov (bias váh vs. žiadny bias váh pre RMS normu, aktivačné funkcie atď.). V dôsledku toho sme sa rozhodli použiť DSL na ich implementáciu: ROSE sme pôvodne vytvorili pomocou Tritonu. Odvtedy sme však ako primárne prostredie pre programovanie GPU prijali CuTeDSL.

Z matematického hľadiska sú inferenčné jadrá zvyčajne nenáročné. V dôsledku toho pri vývoji jadra spočíva zložitosť vo všeobecnosti nie v tom, aký výpočet sa má vykonať, ale v tom, ako presne by ho mal hardvér vykonať. Na vytvorenie najlepších jadier potrebujeme jazyk, ktorý ponúka správne abstrakcie nad základným hardvérom. Správny jazyk nám umožňuje špecifikovať operácie na vysokej úrovni na vykonanie výpoctu, pričom ich agresívne špecializujeme a optimalizujeme vo všetkých bodoch v širokom konfiguračnom priestore.

CuTeDSL je jazyk špecifický pre doménu (DSL) pre GPU NVIDIA založený na algebre rozvrhnutia CuTe a MLIR. Kompiluje sa za chodu (JIT) do vysoko optimalizovaného PTX, pričom si zachováva kontrolu nad nízkoúrovňovými hardvérovými primitívami. Jeho cieľom je ponúknuť rovnaké funkcie a výkon ako CUTLASS C++, bez nákladných časov kompilácie, ktoré brzdia rýchlosť vývoja. Z jazykov, ktoré sme preskúmali alebo s ktorými sme pracovali (vrátane CUDA, cuTile, TVM IR a Triton), ho považujeme za najvhodnejší pre vývoj našich jadier.

Celkovo, hoci nie je taký na vysokej úrovni a stručný ako Triton alebo cuTile, zistili sme, že algebra rozvrhnutia CuTe a jej kompozície (sprístupnené prostredníctvom cute.Tensor) sú rozumnou náhradou. CuTeDSL navyše ponúka neobmedzený prístup k hardvérovým primitívam, čo nám poskytuje úroveň kontroly, akú potrebujeme na dosiahnutie špičkového výkonu vo všetkých okrajových prípadoch. S CUDA boli možnosti špecializácie v čase kompilácie obmedzené, pretože každý ďalší parameter pridával neprimerané náklady na časy kompilácie alebo zložitosť systému zostavovania. Na rozdiel od toho s DSL kompilovaným za chodu (Just-In-Time) môžeme špecializovať agresívnejšie bez toho, aby sme obetovali tempo iterácie. Na rozdiel od TVM IR, ktorý používal CUDA ako medzikrok, sme zistili, že pipeline založená na MLIR, ktorá prechádza na PTX, je robustnejšia a jednoduchšia na ladenie ako tá, ktorá emitovala CUDA ako medzikrok. Okrem toho sme zistili, že vďaka použitiu vlastných operácií na zabalenie inštrukcií PTX môžeme naplno využiť všetky funkcie základného hardvéru.

Prevažná časť nákladov na kompiláciu jadier CUDA sa zvyčajne vynakladá vo fáze rozširovania šablón jadier, keďže optimalizáciu a kompiláciu od fázy PTX ďalej možno dokončiť v priebehu niekoľkých sekúnd. Pri mnohých parametroch v čase kompilácie môže táto fáza kompilácie narásť na niekoľko minút pre netriviálne jadro. Toto úzke hrdlo je umocnené skutočnosťou, že možnosti paralelnej kompilácie vznikajú až neskôr, na úrovni IR a PTX, pričom spracovanie šablón C++ prebieha sekvenčne. Hoci existujú triky na riešenie tohto problému presunutím inštancií šablón do samostatných kompilačných jednotiek, vyžadujú si významné úpravy systémov zostavovania a používanie ťažkopádnych makier. Naproti tomu jadrá CuTeDSL, ktoré sú kompilované za chodu (JIT), nevyžadujú zložité pravidlá zostavovania a nebrzdia časy kompilácie, pretože cyklus vývoja sa môže zamerať na jeden bod v konfiguračnom priestore, zatiaľ čo starostlivé jednotkové testy môžu neskôr pokryť celý návrhový priestor jadra.

Aj keď JIT kompilácia presúva náklady z času vývoja do času bežného chodu (runtime), réžia je viac než prijateľná. Keď sa model nasadí, aktivuje len niekoľko bodov z celého konfiguračného priestoru, čo spôsobuje režijné náklady v ráde sekúnd počas spustenia/zahriatia inferenčného enginu. Naproti tomu predbežná kompilácia všetkých konfigurácií jadra by mohla trvať niekoľko minút. Okrem toho sa réžia dá ďalej zmierniť ukladaním skompilovaných artefaktov do vyrovnávacej pamäte (caching).

Vysoké náklady na kompiláciu CUDA sú umocnené obtiažnosťou sledovania pôvodu chýb, keď sa používa dispečing a špecializácia založená na šablónach, čo je prípad CuTe a CUTLASS. S DSL je možné chyby nasmerovať na riadok kódu, ktorý ich vygeneroval, čo výrazne uľahčuje ladenie. Navyše je oveľa jednoduchšie mapovať vygenerovaný kód na príkaz DSL, z ktorého pochádza, prostredníctvom stredných fáz MLIR alebo spoliehaním sa na ladiace informácie priradené k vygenerovanému PTX. Vďaka týmto aspektom DSL sme zistili, že sme pri vývoji jadier oveľa produktívnejší.

Hoci jadrá zdieľajú veľa podobností naprieč modelmi, stále musia byť vybavené radom konfiguračných ovládacích prvkov, ktoré môžu mať vplyv na výkon. Relevantné parametre zahŕňajú:

Skrytá dimenzia / Dimenzia hlavy / Počet hláv pozornosti

Vstupné/výstupné dátové typy

Schéma kvantovania

Konštanty v čase kompilácie (biases, aktivačné funkcie atď.)

Architektúra ROSE vrátane jej enginu, modelu, knižnice jadra a služieb.

Špecializácia v čase kompilácie je užitočná pre výkon aj ladenie. Najmä v paralelnom procesore, akým je GPU, je neexistujúca vetva lepšia ako vetva, ktorá sa berie vždy alebo nikdy. Okrem toho kompilácia jadra pre konkrétnu súpravu konfiguračných ovládacích prvkov vedie k menšiemu počtu inštrukcií PTX/SASS, ktoré sa ľahšie analyzujú a ladia pri použití cuda-gdb počas vyšetrovania.

Pre danú veľkosť dávky a konfiguráciu je možné odhadnúť očakávanú premávku do DRAM a počet aritmetických operácií. Následne je možné urobiť optimálne rozhodnutia od prípadu k prípadu pre optimálnu alokáciu mriežky (grid allocation), synchronizáciu a využitie určitých funkcií hardvéru. CuTeDSL nám umožňuje definovať a opätovne použiť komponenty na hĺbkovú špecializáciu bez toho, aby to prinášalo neprijateľné režijné náklady v čase kompilácie.

Vo všeobecnosti musí byť výkon inferenčných jadier optimalizovaný pre dva samostatné prípady použitia:

Prefill (Predbežné plnenie): malá veľkosť dávky (málo blokov/chunkov, sekvencií), veľa tokenov na sekvenciu. Priepustnosť sa má maximalizovať pre vklady (embeddings) aj vstupné tokeny modelov LLM. Latencia v desiatkach až stovkách milisekúnd.

Dekódovanie: malé až veľké veľkosti dávok, jeden (bez špekulatívneho dekódovania) až niekoľko tokenov na sekvenciu (špekulácie cieľ-návrh, predikcia viacerých tokenov). Latencia blízka latencii spustenia jadra.

Pre GEMM a pozornosť (attention) beží fáza prefill na veľkom počte tokenov a výkon je viazaný výpočtovou priepustnosťou zariadenia. Dekódovanie zahŕňa tú istú operáciu na malých veľkostiach dávok, pričom časy vykonávania sú ovplyvnené nákladmi na načítanie aktivácií a váh z DRAM do streamovacích multiprocesorov a tenzorových jadier, pričom sú namiesto toho obmedzené šírkou pásma pamäte zariadenia. Aj keď je veľkosť dávky alebo celkový počet tokenov v dávke relatívne jasným ukazovateľom, prah prechodu medzi prefill a dekódovaním je jemnejší, pretože šírka pásma pamäte je tiež určená skrytou dimenziou modelu a dátovým typom aktivácie. Jadrá sú preto rozdelené na implementácie prefill a dekódovania, ak by algoritmické rozdiely alebo priradenie mriežky/bloku zvýhodňovali niektoré z nich.

Parameter, ktorý je univerzálne odlišný medzi jadrami prefill a dekódovania, je počet warpov v bloku (počet vlákien vydelený veľkosťou warpu 32 vlákien na warp). Je žiaduce urobiť tento parameter konštantou v čase kompilácie, pretože môže pomôcť pri alokácii registrov. V CUDA sa to zvyčajne robí špecifikáciou počtu vlákien prostredníctvom __launch_bounds__ a poskytnutím počtu warpov ako parametra šablóny:

cpp

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

}

Počet registrov dostupných pre všetky vlákna v bloku je hardvérové obmedzenie – 64K na Hopper a Blackwell. Ak je počet vlákien známy vopred, kompilátor presne vie, koľko registrov môže priradiť každému vláknu, a môže úspešne kompilovať agresívnejšie optimalizované jadrá. To však prináša náklady v čase kompilácie, pretože teraz je potrebné jadro prekompilovať pre rôzne veľkosti warpov. S CuTeDSL to už nie je problém, pretože nemusíme platiť náklady na kompiláciu pre celý konfiguračný priestor vopred. DSL tiež očakáva, že veľkosť bloku bude konštantným výrazom, čím implicitne vykonáva optimalizáciu implikovanú __latch_bounds__ v CUDA.

Okrem toho, že počet warpov bude konštantný, je žiaduce ho meniť medzi prefill a dekódovaním. Na plné využitie šírky pásma pamäte GPU nestačí vždy priradiť vysoký počet warpov na blok, pretože plnú šírku pásma možno dosiahnuť iba použitím čo najväčšieho počtu dostupných streamovacích multiprocesorov (SM). V dôsledku toho pre daný počet tokenov a známe dátové typy a skryté dimenzie najprv vyberieme ideálnu vektorovú dimenziu na vlákno (8 pre bfloat16, 4 pre float32), aby sme zabezpečili, že vygenerovaný kód používa 128-bitové inštrukcie načítania/uloženia (256-bitové na Blackwell). Keď poznáme veľkosť vektora, môžeme vypočítať počet warpov požadovaných jadrom. Najprv sa snažíme maximalizovať počet SM: ak je menej warpov ako SM, vyberie sa jeden warp na blok. Toto je zvyčajne variant jadra na dekódovanie. Ak je dostatok warpov, postupne zvyšujeme počet warpov na 8, 16 alebo 32, pričom robíme benchmark hornej hranice, po ktorej už zvyčajne nevidíme zlepšenie výkonu.

Zatiaľ čo Triton vykonal niektoré z týchto optimalizácií pod kapotou, zistili sme, že podobné výsledky môžeme dosiahnuť starostlivým použitím algebry rozvrhnutia CuTe bez toho, aby sme boli obmedzení nedostatkami vysokourovňového kompilátora. Najmä vtedy, keď ide o horizontálne redukcie pod warpom, sme zistili, že CuTeDSL nám ponúka jemne zrnitú kontrolu potrebnú na generovanie optimálneho kódu.

V niektorých situáciách, na distribúciu práce medzi čo najviac blokov, možno použiť cooperative_groups::this_grid().sync() na synchronizáciu všetkých zúčastnených blokov. Platia obmedzenia: keďže blok nemožno odplánovať z SM, akonáhle začne vykonávanie, na zaručenie napredovania nesmie počet blokov prekročiť počet SM dostupných na zariadení, aby sa zabezpečilo, že v každom bode môžu všetky bloky aktualizovať a kontrolovať miesta v globálnej pamäti zapojené do synchronizácie. Toto obmedzenie vynucuje cudaLaunchCooperativeKernel. Náklady na mriežkovú bariéru sú približne 3 us na Hopperi plus náklady na vyprázdnenie všetkých predchádzajúcich pamäťových operácií prostredníctvom implicitnej pamäťovej bariéry v niektorých situáciách. Celkovo je to porovnateľné s réžiou rozdelenia jadra na dve spustenia, čo vedie k rozhodnutiu o návrhu citlivému na výkon. Napriek tomu, že programatívne závislé spustenie (PDL) môže znížiť réžiu spustenia od Hoppera vyššie, pre jadrá dekódovania sa oplatí vyhnúť sa tejto réžii.

Ak jedno z pracovných zaťažení vyžaduje viac blokov, než je počet SM, jadro by sa malo rozdeliť, pričom sa spolieha na implicitnú bariéru spustení jadra. To je primárne prípad prefill. Ak sa počet blokov, ktoré sú ideálne pre všetky pracovné zaťaženia, zmestí do mriežky, uprednostňuje sa mriežková bariéra, ako je to v prípade dekódovacích jadier operujúcich s menším počtom tokenov.

V CuTeDSL je možné tieto rozhodnutia urobiť v spúšťacom kóde na strane hostiteľa a jadrá je možné skladať a prekladať mriežkovými bariérami podľa potreby, ako je to znázornené na príklade nižšie. S takouto hybridnou stratégiou môžu byť dekódovacie jadrá optimalizované tak, aby dosiahli minimálne latencie, a predbežné plnenie (prefill) môže byť optimalizované pre maximálnu priepustnosť. Na rozdiel od riešení založených na CUDA možno túto logiku definovať stručnejšie a s menšími režijnými nákladmi v čase kompilácie.

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

V prípade našich MoE all-to-all jadier, ako je podrobne uvedené nižšie, nám flexibilita vytvárať jadrá s mriežkovými bariérami alebo bez nich a optimalizovať ich v čase kompilácie pomohla dosiahnuť výrazné prírastky výkonu prostredníctvom CuTeDSL.

Prvé jadro, ktoré sme vytvorili v CuTeDSL, bolo špecializované jadro RMS norm, optimalizované pre malé dimenzie. QK norm, používaná vo vrstve self-attention nedávnych modelov Qwen, sa aplikuje na základe jednotlivých hláv (per-head) na 64, 128 alebo 256 prvkov. Na rozdiel od bežnej RMS normy, ktorá operuje na skrytých dimenziách presahujúcich tisíce a vyžaduje horizontálnu redukciu cez zdieľanú pamäť aj registre, QK norm vyžaduje redukciu iba naprieč warpmi s možnosťou vtlačiť viacero hláv do jediného warpu. Toto sa stalo prípadom použitia, kde sme mohli agresívne optimalizovať v čase kompilácie, aby sme našli najlepšie priradenia tokenov k warpom pre všetky dimenzie hláv a dátové typy.

RMS norm je pomerne priamočiara operácia, aj keď vyžaduje horizontálny súčet na výpočet štvorcovej normy vektora:

cpp

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

Keďže hlavy pozornosti (attention heads) sú zvyčajne malé a operácie na vlákno sú silne vektorizované, jeden warp môže vypočítať štvorcovú normu. S cieľom maximalizovať šírku pásma pamäte by každé vlákno malo spracovať 128 bitov vstupu, čo sa prekladá na 8 hodnôt bfloat16. Polovičný warp (half-warp) je preto postačujúci na spracovanie hlavy a štvorcovú normu možno určiť pomocou horizontálnej redukcie vykonanej pomocou shuffle_sync_bfly. Vychádzajúc z dimenzie hlavy a dátového typu, môžeme určiť počet hláv na warp a špecifikovať argumenty a počet kôl redukcie v čase kompilácie.

Potom môžeme zostaviť špecializované jadro s pomerne malým počtom riadkov kódu, pričom sa spoliehame na automaticky vektorizované načítania/uloženia a používame operácie špecifické pre architektúru nižšej úrovne na redukciu warpov. Jadro určuje, na ktorej hlave (Q alebo K) sa má pracovať, a identifikuje zodpovedajúce váhy, ktoré sa majú aplikovať. Načíta celý vektor z hlavy a najprv ho lokálne sčíta. Následne sa na základe počtu vlákien na hlavu vypočíta štvorcová norma. Nakoniec sa určí automaticky vektorizovaná RMS norma, ktorá sa zapíše späť do pamäte. Parametre špecifické pre model, ako napríklad WEIGHT_BIAS a EPS, sú zložené z konštánt (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)

V porovnaní s bežným jadrom RMS norm poskytuje tento špecializovaný variant 2 až 3-krát lepšiu priepustnosť a latenciu pre prefill aj dekódovanie.

MoE dispečing a kombinovanie premiešavajú vstupné tokeny odborníkom (experts), na ktorých sú smerované, a vykonávajú váženú akumuláciu po tom, čo jadrá MoE GEMM vykonajú príslušné násobenia matíc. Výkon je predmetom mimoriadneho záujmu, keďže rôzne produkty sa spoliehajú na modely MoE rôznych veľkostí. Na základe typov nasadenia rozlišujeme rôzne jadrá na optimalizáciu:

TP=1: Lokálne jadro pre jediné GPU. Dispečing premiešava tokeny príslušným odborníkom a kombinovanie ich kumuluje. Používa sa primárne pre nasadenia Qwen3 30B na H100.

DP=1, TP ≠ 1: Tenzorovo-paralelné, jednouzlové jadro. Vstup je replikovaný, dispečing vyberá tokeny pre aktuálny rank a kombinovanie ich kumuluje naprieč všetkými rankmi. Akumulácia cez NVLink alebo NVLS. Používa sa pre veľké prefillery.

DP ≠ 1, NVLink: Dátovo-paralelné jadro s voliteľným tenzorovým paralelizmom. Každý rank posiela tokeny každému inému ranku, ktorý hostí špecifických odborníkov. Informácie o smerovaní sú agregované, tokeny sa prenášajú cez NVIDIA NVLink. Používa sa pre dátovo-paralelné nasadenia H200 alebo GB200.

DP ≠ 1, InfiniBand: Peer-to-peer dátovo-paralelné a tenzorovo-paralelné jadro fungujúce naprieč NVLink aj NVIDIA Quantum InfiniBand. Používa sa pre viacuzlové nasadenia Hopper.

Lokálne jadro počítalo počet tokenov priradených každému odborníkovi na základe smerovacej tabuľky a priraďuje unikátny index každému tokenu v rámci zoznamu tokenov patriacich každému odborníkovi. Po kroku smerovania sa tokeny skopírujú do voliteľne zarovnaných (padded) vyrovnávacích pamätí, ktoré sa odovzdajú jadru MoE GEMM. Pomocou umiestnení tokenov potom kombinovacie jadro vykoná akumuláciu v registroch a zapíše finálne tokeny do výstupnej vyrovnávacej pamäte.

Fáza dispečingu sa líši medzi prefill a dekódovaním: pre malý počet tokenov je k dispozícii len málo informácií o smerovaní na prečítanie, takže ich spracujeme v jedinom bloku a signalizujeme zvyšok mriežky, pričom jednotlivé bloky kopírujú bloky tokenov. Priradenie tokenov sa určuje atomickým inkrementovaním počítadiel v zdieľanej pamäti. Keďže prefill musí spracovať viac informácií o smerovaní, smerovanie je rozdelené do viacerých blokov v samostatnom jadre. Indexy sa najprv priradia lokálne pomocou počítadiel v zdieľanej pamäti a potom sa agregujú naprieč všetkými blokmi v globálnej pamäti, aby sa určila pozícia každého tokenu. Mriežka tohto jadra je obmedzená (clamped) na počet SM, aby bolo možné použiť bariéry na úrovni mriežky. Následné jadro vykoná kópie, spúšťané naprieč oveľa väčším počtom blokov a warpov, aby sa naplno využila dostupná šírka pásma pamäte.

Kombinové jadro je pomerne jednoduché, pretože vykonáva vektorizovanú akumuláciu, pričom sa spolieha na informácie o smerovaní a posuny vypočítané dispečingom. CuTeDSL sa tu ukázal ako cenný, pretože výpočet bolo možné vyjadriť pomocou jednoduchých tenzorov CuTe, ktoré sa kompilujú do efektivneho kódu pre všetky dátové typy a skryté dimenzie.

V tenzorovo-paralelnom nasadení sú informácie o smerovaní a vstupné tokeny replikované. Lokálne dispečerské jadro sa dá plne opätovne použiť s miernou zmenou na smerovanie a výber iba tých tokenov, ktoré sú priradené odborníkom umiestneným na aktuálnom ranku. S niektorými parametrami a špecializáciou v čase kompilácie sa lokálne dispečerské jadrá opätovne používajú.

Keďže počet rankov je zvyčajne nižší ako počet odborníkov, na ktorých je token smerovaný, v priemernom prípade možno predpokladať, že každý token má aspoň jednu zo svojich replík smerovanú na každý rank. V tejto situácii možno lokálne kombinové jadro opätovne použiť na agregáciu lokálnych kópií tokenu, po ktorej nasleduje all-reduce na sčítanie naprieč zúčastnenými rankmi. V prípade predbežného plnenia (prefill kernel) spúšťame samostatné jadrá, avšak dekódovanie spája obe operácie do jediného spustenia.

Pre all-reduce používame obyčajný NVLink, ak sa zúčastňujú iba dva ranky, pričom prechádzame na NVLS, keď je zapojených 4 alebo viac rankov, pre 10-20% nárast priepustnosti. Vďaka DSL môžeme zdieľať väčšinu redukčnej logiky a výmen peer vyrovnávacích pamätí medzi MoE a bežnými all-reduce jadrami. Toto je situácia, kde je prospešná úzka kontrola nad hardvérom, pretože operácie DSL nad multimem.ld_reduce a multimem.st sú vyžadované pre NVLS. Samotná implementácia redukcie do veľkej miery nasleduje príklad all_reduce z testov NCCL.

Dátovo paralelné jadrá sa spoliehajú na NVLink na agregáciu tokenov pre viaceré zariadenia. Vo fáze dispečingu, pred výmenou tokenov, rovnocenné uzly (peers) najprv smerujú lokálne, aby identifikovali cieľový posun na každom ranku pre každý token, ktorý odosielajú. Následne začnú kopírovať tokeny do prijímacích vyrovnávacích pamätí alokovaných pre každého odosielateľa, aby sa predišlo konfliktom. Na signalizáciu dokončenia sa používajú zápisy typu release-store na token. Na strane prijímača sú informácie o smerovaní agregované, aby bolo možné skondenzovať všetky tokeny do jedinej vyrovnávacej pamäte pre MoE GEMM. Prijímače kontrolujú vlajku zapísanú peerom pomocou acquire-load, aby počkali na dokončenie zápisov.

Na strane kombinovania sa tokeny kopírujú cez NVLink do zdrojových rankov, pričom vlajka na token označuje prijatie. Kombinové jadro je podobné lokálnym a jednouzlovým variantom s miernymi rozdielmi v adresovaní a s určitou dodatočnou synchronizačnou logikou.

Na implementáciu tohto jadra sme vytvorili knižnicu obalov PTX na sprístupnenie načítaní a uložení typu release-acquire vymedzených na úrovni systému aj zariadenia. Zlepšením granularity synchronizačných vlajok a alokovaním viacerých blokov na odosielanie a prijímanie tokenov cez NVLink sme dosiahli 2-násobné zlepšenie oproti našim predchádzajúcim jadrám CUDA.

Naše peer-to-peer jadrá predtým dosahovali špičkový výkon na GPU Hopper spárených s adaptérmi ConnectX-7 InfiniBand. Odvtedy sme portovali kombinové jadrá na CuTeDSL pre ďalšie zisky vo výkone. Keďže verzie CUDA sa už silne spoliehali na špecializáciu v čase kompilácie, ukázali sa ako úzke hrdlá času kompilácie v našom repozitári, čo nás nútilo vynechať niektoré parametre. Migráciou na JIT kompilované DSL sme špecializovali pre viac parametrov a vyladili alokácie mriežky, čím sme zaviedli ďalšie rozdiely medzi prefill a dekódovaním pre približne 3% zvýšenie latencie.

Z perspektívy jadier ROSE spadá podpora Blackwell do dvoch kategórií: jadrá agnostické voči architektúre a jadrá špecifické pre architektúru. Jadrá špecifické pre architektúru, ktoré sa spoliehajú na tenzorové jadrá, vyžadujú explicitnú prácu na prispôsobenie. Vďaka použitiu DSL sa však jadrá agnostické voči architektúre prispôsobujú bez problémov (triviálne).

GPU Blackwell má až 160 SM, na rozdiel od 132 v prípade Hoppera. Pre väčšinu jadier je to konštanta v čase kompilácie a logika v spúšťacom kóde na strane hostiteľa to berie do úvahy pri určovaní optimálneho jadra (prefill vs. dekódovanie) na spustenie. Blackwell navyše zavádza 256-bitové inštrukcie načítania, ktoré možno jednoducho emitovať zväčšením šírky vektora (posledná dimenzia tenzorov CuTe) na zlepšenie využitia šírky pásma pamäte.

Blackwell zavádza hardvérovú podporu pre významný počet schém kvantovania (MXFP8, MXFP4, NVFP4). Podpora pre ne bola implementovaná primárne vyladením existujúcich jadier pre blokové škálovanie FP8. Podpora JIT je tu kľúčová, pretože nové režimy môžu byť podporované bez explózie konfiguračného priestoru v čase kompilácie. V minulosti sme zistili, že Triton je neohrabaný pri vyjadrovaní horizontálnych redukcií naprieč niekoľkými prvkami, nad polovicou alebo štvrťrokou warpu. Jemnejšia kontrola nad redukciami bola kľúčová na dosiahnutie optimálneho výkonu pri výpočte absolútnej maximálnej hodnoty naprieč 16 alebo 32 prvkami pre NVFP4 a MXFP8/MXFP4.

Ako relatívne nový jazyk mali LLM malú expozíciu voči kódu CuTeDSL počas ich predbežného tréningu. Okrem toho štýl nášho súčasného kódu s jadrami optimalizovanými pre aktivácie a normy bol mierne odlišný od existujúcich príkladov CUTLASS alebo Quack, ktoré boli silne optimalizované pre TMA a inštrukcie násobenia matíc špecifické pre architektúru. Dosiahnutie bodu, kde by LLM mohli pomáhať pri písaní jadier, bol postupný proces.

Manuálne sme vybudovali počiatočný korpus jadier a pomocných funkcií, ktoré sa držali konzistentných vzorov, skôr než sme sa mohli spoliehať na agentov pri konverzii jadier. S napredovaním korpusu sme mohli vykonávať hrubé preklady z CUDA do CuTeDSL, po ktorých nasledovalo významné manuálne prerábanie. Zistili sme najmä, že modely bojovali s konverziami ukazovateľov na tenzory a vektorizáciou a uviazli pri pokusoch o opravu chýb typov, keď sa snažili prinútiť test ku kompilácii. Ukázali len malé pochopenie povahy jadier založenej na mriežke, pričom si mýlili synchronizáciu vlákien a mriežky. Najviac frustrujúca bola snaha halucinovať neexistujúce preklady knižníc špecifických pre CUDA namiesto vytvárania na základe našich vlastných náhrad CuTeDSL. Postupom času, ako korpus jadier CuTeDSL rástol, sme zistili, že LLM dokážu vytvoriť väčšinu kostry kódu a vyplniť prázdne miesta, čo nám umožňuje sústrediť sa na dizajn na vysokej úrovni.

V porovnaní s CUDA zisťujeme, čo sa týka CuTeDSL, že prináša oveľa lepší výkon v čase kompilácie, čo nám umožňuje špecializovať sa na viac parametrov a dosiahnuť zlepšenie výkonu. Hoci sú jadrá o niečo obšírnejšie kvôli absencii určitých abstrakcií v porovnaní s Tritonou, zistili sme, že zložitosť je oprávnená, pretože jemnejšia kontrola nad hardvérom nám umožňuje vytvárať rýchlejšie jadrá.

V čase písania tohto článku je výkonnostná medzera medzi jadrami CUTLASS C++ a CuTeDSL minimálna, čo otvára zaujímavé príležitosti na komponovanie a spájanie vysoko výkonných jadier GEMM s komunikáciou medzi zariadeniami, normalizáciou alebo aktiváciami. Záväzkom k tomuto DSL budeme schopní vydolovať z našich zariadení ďalší výkon na zlepšenie priepustnosti naprieč našimi pracovnými zaťaženiami, od vkladov (embeddings) až po LLM.

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

https://github.com/NVIDIA/cutlass

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

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

BibTeX

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

Inferencia v Perplexity AI

Prečo CuTeDSL?

Vývojárska skúsenosť

Veľký konfiguračný priestor

Predvyplnenie a dekódovanie špecializácie

Synchronizácia mriežky

Jadrá CuTeDSL

QK Norm

Mixture-of-Experts Dispatch / Combine (Odosielanie / Kombinovanie zmesí odborníkov)

Lokálne

Jednouzlové (Singlenode)

Dátový paralelizmus (Data Parallel)

Peer-to-Peer

Podpora pre Blackwell

Generovanie jadier asistované AI

Záver a budúca práca

Referencie

Citácia