CuTeDSL at Perplexity

Perplexity „Runtime-Optimized Serving Engine“ (ROSE) naudojame viduje, kad aptarnautume modelius nuo įdėtinių rodiklių iki trilijonų parametrų LLM. Mūsų vykdymo aplinka, esanti visų API, įskaitant „Sonar“, „Search“ ir „Embeddings“, užkulisiuose, sukurta tam, kad suteiktų tyrėjams galimybę tyrinėti tinkamiausius modelius ir architektūras bei greitai diegti geriausius kandidatus į gamybą. Kad to pasiektume, turime sugebėti sparčiai diegti, pritaikyti ir optimizuoti naujausių architektūrų išvadų darymą. Šiame tinklaraščio įraše nagrinėjame, kaip į mūsų išvadų darymo varikliuką integruoti „CuTeDSL“ branduoliai leidžia mums sukurti reikiamus GPU branduolius, kad modeliai pasiektų didžiausią našumą naudojant „NVIDIA“ GPU.

Pasirinktiniai modeliai, sukurti įvairių „Perplexity“ komandos, yra talpinami viduje „NVIDIA Hopper“ ir „Blackwell“ GPU, naudojant ROSE išvadų darymo varikliuką, kurį sukūrė AI išvadų darymo komanda. ROSE iš pradžių buvo sukurtas siekiant geriau aptarnauti pritaikytus „Llama“ modelius kalbos dekodavimui ir klasifikavimui su sąsaja, suderinama su „NVIDIA Triton Inference Server“. Laikui bėgant jis išaugo į visavertį išvadų darymo varikliuką, aptarnaujantį didelius LLM, taip pat reitingavimo, klasifikavimo, vertinimo ir kitus transformatoriais pagrįstus modelius įvairioms komandoms, kad jos galėtų maitinti „Perplexity“ produktus. ROSE kūrimo tikslas yra būti paprastam: laikui bėgant investuojamos pastangos didžiausiam našumui pasiekti, tačiau pagrindinis prioritetas yra prisitaikyti prie nuolat besikeičiančių mūsų modelių architektūrų ir naudojimo atvejų.

Duomenų tipų, galvutės matmenų, kompiliavimo laiko konstantų ir paleidimo nustatymų apimanti branduolio konfigūracijos erdvė.

Aukštu lygiu ROSE susideda iš varikliuko, kuris pritaiko modelį prie sąsajos, su kuria gali sąveikauti klientai. Varikliuko komponentai yra labai tinkami pakartotiniam naudojimui, tvarko užklausų planavimą ir paketavimą bei gauna dekoduotus tokenus, įdėtinius rodiklius arba įverčius iš esminių modelių. Jie atlieka įrenginių inicializavimą, tarpresesinį ir tarptinklinį ryšio kanalų valdymą bei svorių įkėlimą. LLM atveju varikliukas atsakingas už paketavimą, skaidymą dalimis ir atranką. Jis taip pat atsakingas už KV atminties paskirstymą tiek pilnam, tiek tiesiniam dėmesiui, taip pat priešdėlių atitikimą. Įdėtinių rodiklių varikliukai yra paprastesni, nes jie atlieka tik tiesioginį paketavimą. Vykdymas yra optimizuotas taip, kad varikliukas ruošia kitą darbo ciklą, kol modelis tuo metu lygiagrečiai veikia greitintuve.

ROSE pateikia turtingą pasirinktinių sluoksnių rinkinį modeliavimui, apvyniojantį GPU branduolius, sukurtus su „CuTeDSL“, „Triton“, „CUDA“, „CUTLASS“ ir „cuBLAS“. Matricų daugybai ir dėmesio realizavimui mes daugiausia pasitikime „NVIDIA“ branduoliais iš „CUTLASS“ ir „cuBLAS“. Likusi konvejerio dalis, įskaitant įdėtinių rodiklių, normos, MoE ir aktyvinimo branduolius, yra ROSE branduolių bibliotekos dalis. Siekiant geriausiai palaikyti aparatinę įrangą, kurioje vykdome išvadų darymą, ROSE turi savo tarpįrenginio ryšio primityvus MoE nukreipimui, paskirstymui ir sujungimui. Nors kiti branduoliai yra palyginti paprasti, jie turi būti kruopščiai specializuojami, kad būtų pasiektas didžiausias našumas visuose skirtinguose paslėptuose matmenyse ir įrenginiuose, kuriuose jie veikia. Be to, mes remiamės kompiliavimo laiko specializacija, kad pritaikytume visus įvairių modelių parametrus ir koregavimus (svorio poslinkis ir be svorio poslinkio RMS normai, aktyvinimo funkcijos ir kt.). Dėl to nusprendėme jiems įgyvendinti naudoti DSL: iš pradžių ROSE sukūrėme naudodami „Triton“. Tačiau nuo to laiko „CuTeDSL“ priėmėme kaip pagrindinę GPU programavimo aplinką.

Matematiniu požiūriu išvadų darymo branduoliai dažniausiai yra nesudėtingi. Todėl kuriant branduolį sudėtingumas paprastai slypi ne skaičiavimų išreiškime, o tame, kaip tiksliai aparatinė įranga turėtų juos atlikti. Kad sukurtume geriausius branduolius, mums reikia kalbos, kuri siūlytų tinkamas abstrakcijas esamos aparatinės įrangos atžvilgiu. Tinkama kalba leidžia mums nurodyti aukšto lygio operacijas skaičiavimams atlikti, aktyviai jas specializuojant ir optimizuojant visuose plačios konfigūracijos erdvė taškuose.

„CuTeDSL“ yra „Python“ pagrįsta DSL, skirta „NVIDIA“ GPU, sukurta ant „CuTe“ išdėstymo algebros ir MLIR viršaus. Ji realiuoju laiku (JIT) kompiliuojama į itin optimizuotą PTX, kartu išlaikant valdymą žemo lygio aparatinės įrangos primityvams. Jos tikslas – pasiūlyti tokias pačias funkcijas ir našumą kaip ir „CUTLASS C++“, be brangiai kainuojančio kompiliavimo laiko, kuris stabdo kūrimo greitį. Iš kalbų, kurias ištyrėme arba su kuriomis dirbome (įskaitant „CUDA“, „cuTile“, „TVM IR“ ir „Triton“), manome, kad ji labiausiai tinka mūsų branduolių kūrimui.

Apskritai, nors ir ne tokia aukšto lygio bei glausta kaip „Triton“ ar „cuTile“, nustatėme, kad „CuTe“ išdėstymo algebra ir jos kompozicijos (pateikiamos per cute.Tensor) buvo pagrįsti pakaitalai. Be to, „CuTeDSL“ suteikia neribotą prieigą prie aparatinės įrangos primityvų, suteikdama mums tokį valdymo lygį, kurio reikia didžiausiam našumui pasiekti visais kraštutiniais atvejais. Naudojant „CUDA“, kompiliavimo laiko specializacijos galimybės buvo ribotos, nes kiekvienas papildomas parametras padidino nepagrįstą kompiliavimo laiko sąnaudų ar kompiliavimo sistemos sudėtingumo didėjimą. Priešingai, naudojant realiuoju laiku kompiliuojamą DSL, galime specializuotis agresyviau neaukodami iteravimo tempo. Skirtingai nuo „TVM IR“, kuri kaip tarpinį žingsnį naudojo „CUDA“, nustatėmė, kad MLIR pagrįstas konvejeris, nuleidžiantis į PTX, yra tvirtesnis ir lengviau derinamas nei tas, kuris išvedė „CUDA“ kaip tarpinį žingsnį. Be to, naudodami pasirinktines operacijas PTX instrukcijoms apvynioti, nustatėme, kad galime visiškai išnaudoti visas esamos aparatinės įrangos funkcijas.

Didžioji „CUDA“ branduolių kompiliavimo sąnaudų dalis paprastai tenka šablonų išplėtimo fazei, nes optimizavimas ir kompiliavimas nuo PTX stadijos gali būti atlikti per kelias sekundes. Turint daug kompiliavimo laiko parametrų, šis kompiliavimo etapus gali išaugti iki kelių minučių nesudėtingam branduoliui. Šis butelio kaklelis dar labiau pabrėžiamas tuo faktu, kad lygiagretio kompiliavimo galimybės atsiranda tik vėliau, IR ir PTX lygiais, o C++ šablonų tvarkymas vyksta nuosekliai. Nors egzistuoja triukai tam išspręsti, perkeliant šablonų egzempliorius į atskirus kompiliavimo vienetus, jiems reikia reikšmingų kompiliavimo sistemų pataisymų ir nerangių makrokomandų naudojimo. Priešingai, realiuoju laiku (JIT) kompiuojami „CuTeDSL“ branduoliai nereikalauja sudėtingų kompiliavimo taisyklių ir netrukdo kompiliavimo laikui, nes kūrimo ciklą galima sutelkti į vieną konfigūracijos erdvė tašką, o kruopštūs vieneto testai vėliau gali padengti visą branduolio dizaino erdvę.

Net jei JIT kompiliavimas perkelia sąnaudas iš kūrimo laiko į vykdymo laiką, pridėtinės išlaidos yra daugiau nei priimtinos. Kai modelis yra diegiamas, jis aktyvuoja tik kelis taškus iš visos konfigūracijos erdvė, todėl išvadų darymo varikliuko paleidimo / įšildymo metu patiriamos sekundžių eilės pridėtinės išlaidos. Priešingai, iš anksto kompiuojant visas branduolio konfigūracijas, tai galėtų užtrukti kelias minutes. Be to, pridėtines išlaidas galima dar labiau sušvelninti talpinant JIT kompiuotus artefaktas talpykloje.

Dideles „CUDA“ kompiliavimo sąnaudas dar labiau padidina sunku atsekti klaidų kilmę, kai naudojamas šablonais pagrįstas paskirstymas ir specializacija, kaip yra „CuTe“ ir „CUTLASS“ atveju. Naudojant DSL, klaidas galima nukreipti į kodo eilutę, kuri jas sugeneravo, taip ženkliai palengvinant derinimą. Be to, daug paprasčiau suderinti sugeneruotą kodą su DSL teiginiu, kuris jį sukūrė per tarpinius MLIR etapus arba remiantis derinimo informacija, pridėta prie sugeneruoto PTX. Šių DSL aspektų dėka pastebėjome, kad kurdami branduolius tapome kur kas produktyvesni.

Nors branduoliai turi daug panašumų tarp modelių, jie vis tiek turi būti aprūpinti įvairiais konfigūracijos nustatymais, kurie gali turėti įtakos našumui. Susiję parametrai apima:

Paslėptas matmuo / galvutės matmuo / dėmesio galvučių skaičius

Įvesties / išvesties duomenų tipai

Kvantavimo schema

Kompiliavimo laiko konstantos (poslinkiai, aktyvinimo funkcijos ir kt.)

ROSE architektūra, įskaitant jos variklį, modelį, branduolio biblioteką ir paslaugas.

Kompiliavimo laiko specializacija yra naudinga tiek našumui, tiek derinimui. Ypač konvejerių procesoriuje, kaip yra GPU atveju, nesamos šakos geriau nei visada paimta ar niekada nepaimta šaka. Be to, sukompiliavus branduolį tam tikram konfigūracijos nustatymų rinkiniui, gaunama mažiau PTX/SASS instrukcijų, kurias lengviau analizuoti ir derinti naudojant cuda-gdb tyrimo metu.

Tam tikram paketo dydžiui ir konfigūracijai galima įvertinti numatomą srautą į DRAM ir aritmetinių operacijų skaičių. Vėliau galima priimti optimalius sprendimus kiekvienu atskiru atveju dėl optimalaus tinklelio paskirstymo, sinchronizavimo ir tam tikrų aparatinės įrangos funkcijų naudojimo. „CuTeDSL“ leidžia mums apibrėžti ir pakartotinai naudoti gilaus specializavimo komponentus nepatiriant nepriimtinų kompiliavimo laiko pridėtinių išlaidų.

Paprastai išvadų darymo branduolių našumas turi būti optimizuotas dviejuose atskiruose naudojimo atvejuose:

Pildymas: mažas paketo dydis (keli gabalai, sekos), daug tokenų sekoje. Pralaidumas turi būti maksimizuojamas tiek įdėtiniams rodikliams, tiek LLM įvesties tokenams. Vėlavimas nuo dešimčių iki šimtų milisekundžių.

Dekodavimas: nuo mažų iki didelių paketų dydžiai, nuo vieno (be spėliojamo dekodavimo) iki kelių tokenų sekoje (projekto taikinio spėliojimas, kelių tokenų prognozė). Vėlavimas artimas branduolio paleidimo vėlavimui.

GEMM ir dėmesio atveju pildymo etapas veikia su dideliu tokenų skaičiumi, o našumas apribojamas įrenginio skaičiavimo pralaidumu. Dekodavimas apima tą pačią operaciją su mažais paketų dydžiais, o vykdymo laikui dominuoja aktyvinimų ir svorių įkėlimo iš DRAM į srauto daugiaprocesorius ir tenzoriaus branduolius sąnaudos, o tai vietoj to ribojama įrenginio atminties pralaidumu. Nors paketo dydis arba bendras tokenų skaičius pakete yra santykinai aiškus rodiklis, perjungimo slenkstis tarp pildymo ir dekodavimo yra subtilesnis, nes atminties pralaidumą taip pat lemia modelio paslėptas matmuo ir aktyvinimo duomenų tipas. Dėl to branduoliai yra padalijami į pildymo ir dekodavimo realizavimus, jei algoritminiai skirtumai arba tinklelio / bloko priskirimai suteiktų pranašumą kuriam nors iš jų.

Parametras, kuris visuotinai skiriasi tarp pildymo ir dekodavimo branduolių, yra gijų skaičius bloke (gijų skaičius, padalintas iš gijų grupės dydžio, kuris yra 32 gijos vienai gijų grupei). Pageidautina padaryti šį parametrą kompiliavimo laiko konstanta, nes tai gali padėti paskirstyti registrus. „CUDA“ tai paprastai atliekama nurodant gijų skaičių per __launch_bounds__ ir pateikiant gijų grupių skaičių kaip šablono parametrą:

cpp

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

}

Registrų, prieinamų visoms bloko gijoms, skaičius yra aparatinės įrangos apribojimas: 64K „Hopper“ ir „Blackwell“ įrenginiuose. Jei gijų skaičius žinomas iš anksto, kompiliatorius tiksliai žino, kiek registrų jis gali priskirti kiekvienai gijai, ir gali sėkmingai kompiliuoti labiau optimizuotus branduolius. Tačiau tai kainuoja kompiliavimo laiko sąnaudas, nes dabar branduolys turi būti perkompiliuotas skirtingiems gijų grupių dydžiams. Naudojant „CuTeDSL“, tai nebėra problema, nes mums nereikia iš anksto mokėti kompiliavimo sąnaudų už visą konfigūracijos erdvė. DSL taip pat tikisi, kad bloko dydis bus pastovi išraiška, netiesiogiai atliekanti optimizavimą, kurį numato __latch_bounds__ „CUDA“.

Be gijų grupių skaičiaus padarymo konstanta, pageidautina jį keisti tarp pildymo ir dekodavimo. Norint visiškai išnaudoti GPU atminties pralaidumą, nepakanka visada priskirti didelio gijų grupių skaičiaus vienam blokui, nes pilną pralaidumą galima pasiekti tik naudojant kuo daugiau turimų srauto daugiaprocesorių (SM). Todėl tam tikram tokenų skaičiui ir žinomiems duomenų tipams bei paslėptiems matmenims pirmiausia pasirenkame idealų vektoriaus matmenį vienai gijai (8 skirta bfloat16, 4 skirta float32), kad užtikrintume, jog sugeneruotame kode naudojamos 128 bitų įkėlimo / saugojimo instrukcijos (256 bitų „Blackwell“). Žinodami vektoriaus dydį, galime apskaičiuoti branduoliui reikalingą gijų grupių skaičių. Pirmiausia stengiamės maksimizuoti SM skaičių: jei gijų grupių yra mažiau nei SM, pasirinkta viena gijų grupė viename bloke. Tai paprastai yra dekodavimo branduolio variantas. Jei gijų grupių pakanka, palaipsniui padidiname gijų grupių skaičių iki 8, 16 arba 32, atlikdami lyginamąjį vertinimą iki viršutinės ribos, po kurios paprastai nebestebime našumo pagerėjimo.

Nors „Triton“ atliko kai kuriuos iš šių optimizavimų

Kai kuriose situacijose, norint paskirstyti darbą kuo didesniam bloku skaičiui, galima naudoti cooperative_groups::this_grid().sync() visiems dalyvaujantiems blokams sinchronizuoti. Taikomi apribojimai: kadangi bloko negalima išregistruoti iš SM, kai jis pradeda vykdymą, siekiant garantuoti pažangą į priekį, blokų skaičius negali viršyti įrenginyje turimų SM skaičiaus, kad būtų užtikrintas, jog bet kuriuo metu visi blokai gali atnaujinti ir tikrinti pasaulinės atminties vietas, susijusias su sinchronizavimu. Šį apribojimą užtikrina cudaLaunchCooperativeKernel. Tinklelio barjero kaina „Hopper“ yra apie 3 us, pridėjus visų ankstesnių atminties operacijų išvalymo kainą per netiesioginį atminties barjerą tam tikrose situacijose. Apskritai tai prilygsta branduolio padalijimo į du paleidimus pridėtinėms išlaidoms, kas lemia našumui jautrų dizaino sprendimą. Nepaisant to, kad programinio priklausomo paleidimo (PDL) funkcija gali sumažinti paleidimo pridėtines išlaidas nuo „Hopper“ įrenginių, dekodavimo branduoliams verta išvengti šių papildomų sąnaudų.

Jei vienai iš darbo apkrovų reikia daugiau blokų nei SM skaičius, branduolys turėtų būti padalintas, pasikliaujant netiesioginiu branduolio paleidimų barjeru. Tai pirmiausia taikoma pildymui. Jei blokų skaičius, kuris yra idealus visoms darbo apkrovoms, telpa tinklelyje, pirmenybė teikiama tinklelio barjerui, kaip yra dekodavimo branduolių, veikiančių su mažesniu tokenų skaičiumi, atveju.

„CuTeDSL“ šie sprendimai gali būti priimti pagrindinio kompiuterio pusės paleidimo kode, o branduoliai gali būti sudaryti ir peripinti tinklelio barjerais pagal poreikį, kaip parodyta žemiau esančiame pavyzdyje. Naudojant tokią hibridinę strategiją, dekodavimo branduoliai gali būti optimizuoti minimaliam vėlavimui pasiekti, o pildymo branduoliai gali būti optimizuoti maksimaliam pralaidumui. Skirtingai nuo „CUDA“ pagrįstų sprendimų, šią logiką galima apibrėžti glaustiau ir su mažesnėmis kompiliavimo laiko pridėtinėmis išlaidomis.

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

Mūsų MoE „viskas su visais“ (all-to-all) branduolių atveju, kaip aprašyta žemiau, lankstumas kurti branduolius su tinklelio barjerais arba be jų ir optimizuoti juos kompiliavimo metu padėjo mums pasiekti didelį našumo prieaugį per „CuTeDSL“.

Pirmasis branduolys, kurį sukūrėme „CuTeDSL“, buvo specializuotas RMS normos branduolys, optimizuotas mažiems matmenims. QK norma, naudojama naujausių „Qwen“ modelių savaiminio dėmesio (self-attention) sluoksnyje, taikoma kiekvienos galvutės pagrindu, 64, 128 arba 256 elementams. Skirtingai nuo įprastos RMS normos, kuri veikia su paslėptais matmenimis, viršijančiais tūkstančius, ir reikalauja horizontalaus sumažinimo per bendrą atmintį ir registrus, QK normai reikia sumažinimo tik tarp gijų grupių, suteikiant galimybę suspausti kelias galvutes į vieną gijų grupę. Tai tapo naudojimo atveju, kai galėjome agresyviai optimizuoti kompiliavimo metu, kad rastume geriausius tokenų priskirimus gijų grupėms visiems galvučių matmenims ir duomenų tipams.

RMS norma yra gana nesudėtinga operacija, nors jai reikia horizontalaus sumavimo vektoriaus kvadratinei normai apskaičiuoti:

cpp

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

Kadangi dėmesio galvutės paprastai yra mažos, o gijų operacijos yra stipriai vektorizuotos, viena gijų grupė gali apskaičiuoti kvadratinę normą. Siekiant maksimizuoti atminties pralaidumą, kiekviena gija turėtų apdoroti 128 bitų įvestį, o tai atitinka 8 bfloat16 reikšmes. Dėl to pusei gijų grupės pakanka galvutei apdoroti, o kvadratinė norma gali būti nustatyta naudojant horizontalų sumažinimą, atliekamą shuffle_sync_bfly. Pradėdami nuo galvutės matmens ir duomenų tipo, galime nustatyti galvučių skaičių vienai gijų grupei ir kompiliavimo metu nurodyti argumentus bei sumažinimo ciklų skaičių.

Tada galime sukurti specializuotą branduolį su palyginti nedaug kodo eilučių, remdamiesi automatiškai vektorizuotais įkėlimais / saugojimais ir naudodami žemesnio lygio architektūrai specifines operacijas gijų grupių sumažinimui. Branduolys nustato, su kuria galvute (Q arba K) dirbti, ir identifikuoja atitinkamus svorius taikymui. Jis įkelia visą vektorių iš galvutės ir pirmiausia susumuoja vietoje. Vėliau, remiantis gijų skaičiumi vienai galvutei, apskaičiuojama kvadratinė norma. Galiausiai automatiškai vektorizuota RMS norma nustato ir įrašoma atgal į atmintį. Modeliui specifiniai parametrai, tokie kaip WEIGHT_BIAS ir EPS, yra sulankstomi konstantomis.

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)

Palyginti su įprastu RMS normos branduoliu, šis specializuotas variantas užtikrina 2–3 kartus geresnį pralaidumą ir vėlavimą tiek pildymo, tiek dekodavimo metu.

MoE paskirstymo ir sujungimo operacijos sumaišo įvesties tokenus su ekspertais, į kuriuos jie nukreipiami, atliekant svertinį kaupimą po to, kai MoE GEMM branduoliai atlieka atitinkamas matricų daugybos operacijas. Naudojimas yra ypač įdomus, nes įvairūs produktai remiasi įvairaus dydžio MoE modeliais. Remdamiesi diegimo tipais, išskiriame skirtingus branduolius optimizavimui:

TP=1: Vietinis branduolys vienam GPU. Paskirstymas sumaišo tokenus su atitinkamais ekspertais, o sujungimas juos kaupia. Naudojama daugiausia „Qwen3 30B“ diegimams H100 įrenginiuose.

DP=1, TP ≠ 1: Tenzorinis lygiagretusis, vieno mazgo branduolys. Įvestis yra kopijuojama, paskirstymas atrenka tokenus dabartiniam rangui, o sujungimas juos kaupia visuose rankuose. Kaupimas per NVLink arba NVLS. Naudojama dideliems pildytuvams.

DP ≠ 1, NVLink: Duomenų lygiagretus branduolys su pasirinktiniu tenzoriniu lygiagretumu. Kiekvienas rangas siunčia tokenus kiekvienam kitam rangui, talpinančiam specifinius ekspertus. Nukreipimo informacija yra agreguojama, tokenai perkeliami per „NVIDIA NVLink“. Naudojama duomenų lygiagretizavimo H200 arba GB200 diegimams.

DP ≠ 1, „InfiniBand“: Tarpusavio (Peer-to-Peer) duomenų lygiagretus ir tenzorinis lygiagretus branduolys, veikiantis per NVLink ir „NVIDIA Quantum InfiniBand“. Naudojama kelių mazgų „Hopper“ diegimams.

Vietinis branduolys suskaičiuoja tokenų, priskirtų kiekvienam ekspertui, skaičių pagal nukreipimo lentelę ir priskiria unikalų indeksą kiekvienam tokenui tokenų sąraše, priklausančiame kiekvienam ekspertui. Po nukreipimo žingsnio tokenai nukopijuojami į pasirinktinai papildytus buferius, kad būtų perduoti į MoE GEMM branduolį. Naudodamas tokenų vietas, sujungimo branduolys atlieka kaupimą registruose, įrašydamas galutinius tokenus į išvesties buferį.

Paskirstymo fazė skiriasi tarp pildymo ir dekodavimo: esant nedideliam tokenų skaičiui, perskaityti nukreipimo informacijos mažai, todėl ją apdorojame viename bloke ir pranešame apie likusią tinklelio dalį, o atskiri blokai nukopijuoja tokenų gabalus. Tokeno priskirimas nustatomas atomiškai didinant skaitiklius bendrojoje atmintyje. Kadangi pildymas turi apdoroti daugiau nukreipimo informacijos, nukreipimas padalijamas į kelis blokus atskirame branduolyje. Indeksai pirmiausia priskiriami vietoje naudojant skaitiklius bendrojoje atmintyje, o po to agreguojami visuose blokuose pasaulinėje atmintyje, kad būtų nustatyta kiekvieno tokeno pozicija. Šio branduolio tinklelis apribojamas iki SM skaičiaus, kad būtų galima naudoti tinklelio lygio barjerus. Vėlesnis branduolys atlieka kopijavimus, paleistus per daug didesnį blokų ir gijų grupių skaičių, kad pilnai išnaudotų turimą atminties pralaidumą.

Sujungimo branduolys yra palyginti paprastas, nes jis atlieka vektorizuotą kaupimą, remdamasis nukreipimo informacija ir poslinkiais, kuriuos apskaičiavo paskirstymas. „CuTeDSL“ čia pasirodė esąs vertingas, nes skaičiavimą galima išreikšti per paprastus „CuTe“ tenzorius, kurie kompiliuojasi į efektyvų kodą visiems duomenų tipams ir paslėptiems matmenims.

Tenzorinio lygiagretumo diegime nukreipimo informacija ir įvesties tokenai yra kopijuojami. Vietinis paskirstymo branduolys gali būti visiškai panaudotas pakartotinai su nedideliu pakeitimu nukreipti ir pasirinkti tik tuos tokenus, kurie priskirti ekspertams, esantiems dabartiniame ranke. Su tam tikrais parametrais ir kompiliavimo laiko specializacija vietiniai paskirstymo branduoliai naudojami pakartotinai.

Kadangi rankų skaičius paprastai yra mažesnis nei ekspertų, į kuriuos nukreipiamas tokenas, skaičius, vidutiniu atveju galima daryti prielaidą, kad kiekvienas tokenas turi bent vieną iš savo kopijų, nukreiptų į kiekvieną ranką. Šioje situacijoje vietinis sujungimo branduolys gali būti pakartotinai panaudotas vietinėms tokeno kopijoms agreguoti, po kurio atliekamas visiškas sumažinimas (all-reduce), kad būtų susumuota visuose dalyvaujančiuose rankuose. Pildymo branduolio atveju paleidžiame atskirus branduolius, tačiau dekodavimo atveju abi operacijos sujungiamos į vieną paleidimą.

Visuotinio sumažinimo (all-reduce) operacijai naudojame paprastą NVLink, jei dalyvauja tik du rankai, perjungdami į NVLS, kai įtraukiami 4 ar daugiau rankų, kad pralaidumas padidėtų 10–20 %. DSL dėka galime bendrinti didžiąją dalį sumažinimo logikos ir tarpusavio buferio mainų tarp MoE ir įprastų visiško sumažinimo branduolių. Tai situacija, kai glaudus aparatūros valdymas yra naudingas, nes NVLS reikalingos DSL operacijos per multimem.ld_reduce ir multimem.st. Paties sumažinimo realizavimas iš esmės seka NCCL testų all_reduce pavyzdį.

Duomenų lygiagretūs branduoliai remiasi NVLink, kad agreguotų tokenus keliems įrenginiams. Paskirstymo fazėje, prieš keičiantis tokenais, partneriai pirmiausia nukreipia vietoje, kad nustatytų tikslinį poslinkį kiekviename ranke kiekvienam siunčiamam tokenui. Vėliau jie pradeda kopijuoti tokenus į gavimo buferius, paskirtus kiekvienam siuntėjui, kad išvengtų konfliktų. Norint pranešti apie užbaigimą, naudojami atskiri tokeno paleidimo-saugojimo (release-store) įrašai. Gavėjo pusėje nukreipimo informacija yra agreguojama, kad visi tokenai būtų sutelkti į vieną buferį MoE GEMM operacijai. Gavėjai tikrina partnerio įrašytą vėliavėlę su įsigijimo-įkėlimo (acquire-load) operacija, kad palauktų, kol įrašai bus baigti.

Sujungimo pusėje tokenai nukopijuojami per NVLink į šaltinio rankus, o vėliavėlė kiekvienam tokenui nurodo gavimą. Sujungimo branduolys yra panašus į vietinio ir vieno mazgo variantus, su nedideliais adresavimo skirtumais ir papildoma sinchronizavimo logika.

Šiam branduoliui įgyvendinti sukūrėme PTX apvalkalų biblioteką, kad atskleistume paleidimo-įsigijimo įkėlimus ir saugojimus, apimančius tiek sistemos, tiek įrenginio lygį. Pagerinę sinchronizavimo vėliavėlių granuliarumą ir paskirstę daugiau blokų tokenams siųsti bei gauti per NVLink, pasiekėme 2 kartus geresnius rezultatus nei ankstesni mūsų „CUDA“ branduoliai.

Mūsų tarpusavio (Peer-to-Peer) branduoliai anksčiau pasiekė pažangiausią našumą „Hopper“ GPU, suporuotuose su „ConnectX-7 InfiniBand“ adapteriais. Nuo to laiko sujungimo branduolius perkėlėme į „CuTeDSL“, kad pasiektume papildomo našumo prieaugio. Kadangi „CUDA“ versijos jau stipriai rėmėsi kompiliavimo laiko specializacija, jos pasirodė esančios kompiliavimo laiko kliūtys mūsų saugykloje, todėl buvome priversti praleisti kai kuriuos parametrus. Migruodami į JIT kompiuojamą DSL, specializavomės daugiau parametrų ir pakoregavome tinklelio paskirstymus, įvesdami tolimesnius skirtumus tarp pildymo ir dekodavimo maždaug 3 % vėlavimo pagerėjimui.

ROSE branduolių požiūriu „Blackwell“ palaikymas skirstomas į dvi kategorijas: architektūrai agnosticinius ir architektūrai specifinius branduolius. Architektūrai specifiniams branduoliams, kurie remiasi tenzoriaus branduoliais, reikia aiškaus darbo pritaikant. Tačiau naudojant DSL, architektūrai agnosticiniai branduoliai prisitaiko be didesnių pastangų.

„Blackwell“ GPU turi iki 160 SM, palyginti su „Hopper“ 132 SM. Daugumai branduolių tai yra kompiliavimo laiko konstanta, o pagrindinio kompiuterio paleidimo kodo logika ją įvertina nustatydama optimalų (pildymo ar dekodavimo) branduolį vykdymui. Be to, „Blackwell“ taip pat pristato 256 bitų įkėlimo instrukcijas, kurias galima lengvai išvesti padidinant vektoriaus plotį (paskutinį „CuTe“ tenzorių matmenį), siekiant pagerinti atminties pralaidumo išnaudojimą.

„Blackwell“ pristato aparatūros palaikymą nemažam skaičiui kvantavimo schemų (MXFP8, MXFP4, NVFP4). Jų palaikymas buvo įgyvendintas daugiausia pakoregavus esamus branduolius FP8 blokų mastelio keitimui. JIT palaikymas čia yra labai svarbus, nes galima palaikyti naujus režimus neišpučiant kompiliavimo laiko konfigūracijos erdvė. Anksčiau nustatėme, kad „Triton“ keistai išreiškia horizontalius sumažinimus per kelis elementus, per pusę arba ketvirtadalį gijų grupės. Smulkesnis sumažinimo valdymas buvo labai svarbus optimaliam našumui pasiekti skaičiuojant absoliučią didžiausią vertę per 16 arba 32 elementus atitinkamai NVFP4 ir MXFP8/MXFP4 atvejais.

Kaip santykinai naujai kalbai, LLM turėjo mažai sąlyčio su „CuTeDSL“ kodu savo išankstinio mokymo metu. Be to, dabartinio mūsų kodo stilius, kai branduoliai optimizuoti aktyvinimams ir normoms, šiek tiek skyrėsi nuo esamų „CUTLASS“ ar „Quack“ pavyzdžių, kurie buvo stipriai optimizuoti TMA ir architektūrai specifinėms matricų daugybos instrukcijoms. Pasiekti tašką, kuriame LLM galėtų padėti rašyti branduolius, buvo palaikymui procesas.

Rankiniu būdu sukūrėme pradinį branduolių ir pagalbinių funkcijų rinkinį, vadovaudamiesi nuosekliais modeliais, prieš galėdami pasikliauti agentais branduolių konvertavimui. Augant rinkiniui, galėjome atlikti apytikslius vertimus iš „CUDA“ į „CuTeDSL“, po kurių seko didelis rankinis perdarymas. Ypač pastebėjome, kad modeliai sunkiai tvarkėsi su žymeklių konvertavimu į tenzorius ir vektorizavimu, įstrigę bandydami ištaisyti tipo klaidas, kai siekė sukompiliuoti testą. Jie parodė mažai supratimo apie tinkleliu pagrįstą branduolių prigimtį, painiodami gijų ir tinklelio sinchronizavimą. Labiausiai nuvylė tai, kad kilo pagunda haliucinuoti neegzistuojančius „CUDA“ specifikavimo bibliotekų vertimus vietoj to, kad būtų kuriama remiantis mūsų pačių „CuTeDSL“ pakaitalais. Laikui bėgant, augant „CuTeDSL“ branduolių rinkiniui, pastebėjome, kad LLM gali sukurti didžiąją kodo karkaso dalį ir užpildyti trūkstamas vietas, leisdami mums sutelkti dėmesį į aukšto lygio dizainą.

Palyginti su „CUDA“, pastebėjome, kad „CuTeDSL“ suteikia daug geresnį kompiliavimo laiko našumą, leidžiantį mums specializuotis daugiau parametrų ir pasiekti našumo pagerėjimą. Nors branduoliai yra šiek tiek žodieniškesni dėl tam tikrų abstrakcijų trūkumo palyginti su „Triton“, manome, kad šis sudėtingumas yra pagrįstas, nes smulkesnis aparatinės įrangos valdymas leidžia mums kurti greitesnius branduolius.

Rašymo metu našumo atotrūkis tarp „CUTLASS C++“ ir „CuTeDSL“ branduolių yra minimalus, o tai atveria įdomias galimybes kurti ir sulieti didelio našumo GEMM branduolius su tarpįrenginio ryšiu, normalizavimu ar aktyvinimais. Įsipareigodami šiam DSL, galėsime išgauti papildomą našumą iš savo įrenginių, kad pagerintume pralaidumą visose mūsų darbo apkrovose, nuo įdėtinių rodiklių iki 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}
}

Išvadų darymas „Perplexity AI“

Kodėl „CuTeDSL“?

Kūrėjo patirtis

Didelė konfigūracijos erdvė

Išankstinis užpildymas ir dekodavimo specializacija

Tinklelio sinchronizavimas

„CuTeDSL“ branduoliai

QK norma

Mixture-of-Experts paskirstymas / sujungimas

Vietinis

Vieno mazgo

Duomenų lygiagretumas

Tarpusavio (Peer-to-Peer)

„Blackwell“ palaikymas

AI pagalba generuojami branduoliai

Išvados ir būsimas darbas

Nuorodos

Citata