CuTeDSL at Perplexity

Perplexitys kasutame me omatehtud käituskeskkonna optimeeritud teenindusmootorit (Runtime-Optimized Serving Engine ehk ROSE), et käitada mudeleid alates manustustest kuni triljoni parameetriga keelemudeliteni (LLM). Kõikide API-de, sealhulgas Sonari, otsingu ja manustuste taustal töötav käituskeskkond loodi selleks, et anda teadlastele võimalus uurida kõige sobivamaid mudeleid ja arhitektuure ning juurutada parimad kandidaadid kiiresti tootmisse. Selle saavutamiseks peame suutma kiiremas korras rakendada, juurutada ja optimeerida järeldustöötlust tipptasemel arhitektuuride jaoks. Selles blogipostituses uurime, kuidas meie järeldusmootorisse integreeritud CuTeDSL-i tuumad (kernelid) võimaldavad meil luua vajalikke GPU tuumasid, et viia mudelid NVIDIA GPU-del maksimaalse jõudluseni.

Kohandatud mudelid, mille on loonud Perplexity erinevad meeskonnad, hostitakse kohapeal NVIDIA Hopperi ja Blackwelli GPU-del, kasutades AI järeldusmeeskonna ehitatud ROSE järeldusmootorit. ROSE loodi algselt selleks, et pakkuda paremat teenust kohandatud Llama mudelitele nii keele dekodeerimiseks kui ka klassifitseerimiseks, kusjuures liides ühildub NVIDIA Triton Inference Serveriga. Aja jooksul kasvas sellest välja täieulatuslik järeldusmootor, mis teenindab märkimisväärseid keelemudeleid, samuti järjestus-, klassifitseerimis-, hindamis- ja muid transfereeripõhiseid mudeleid paljudele meeskondadele, et toetada Perplexity tooteid. ROSE disainieesmärk on olla lihtne: aja jooksul investeeritakse pingutusi tippjõudluse saavutamiseks, kuid peamine prioriteet on kohaneda meie mudelite pidevalt muutuvate arhitektuuride ja kasutusjuhtudega.

Tuuma konfiguratsiooniruum üle andmetüüpide, pea mõõtmete, kompileerimisaja konstantide ja käivitussätete.

Üldises plaanis koosneb ROSE mootorist, mis kohandab mudeli liideseks, millega kliendid saavad suhelda. Mootorid on suurel määral taaskasutatavad, hallates päringute ajastamist ja pakendamist ning hankides alusmudelitelt dekodeeritud tokeneid, manuseid või skoore. Nende ülesandeks on seadmete initsialiseerimine, protsessi- ja sõlmevaheliste sidekanalite seadistamine ning kaalude laadimine. Keelemudelite puhul vastutab mootor pakendamise, tükeldamise ja valimi võtmise eest. Samuti vastutab see KV-mälu eraldamise eest nii täieliku kui ka lineaarse tähelepanu puhul ning eesliite sobitamise eest. Manustusmootorid on lihtsamad, kuna need teostavad ainult veebipõhist pakendamist. Täitmine on optimeeritud nii, et mootor valmistab ette järgmise töövooru samal ajal, kui mudel töötab kiirendil paralleelselt.

ROSE pakub rikkalikku kogumit kohandatud kihte modelleerimiseks, mis ümbritsevad CuTeDSL-i, Tritoni, CUDA, CUTLASS-i ja cuBLAS-i abil ehitatud GPU tuumasid. Maatrikskorrutuse ja tähelepanu rakendamiseks tugineeme peamiselt CUTLASS-i ja cuBLAS-i NVIDIA tuumadele. Ülejäänud torujuhe, sealhulgas manustamis-, normi-, MoE- ja aktiveerimistuumad, kuuluvad ROSE tuumade teeki. Et parimal viisil toetada riistvara, millel me järeldustöötlust teostame, on ROSE varustatud oma seadmetevaheliste sidepõhifunktsioonidega MoE suunamiseks, saatmiseks ja kombineerimiseks. Kuigi teised tuumad on suhteliselt lihtsad, peavad need olema põhjalikult spetsialiseeritud, et saavutada tippjõudlus kõigil erinevatel peidetud mõõtmetel ja seadmetel, millel nad töötavad. Lisaks toetume kompileerimisaegsele spetsialiseerumisele, et kohandada erinevate mudelite kõiki seadistusi ja nüansse (kaalude nihked vs. ilma kaalude niheta RMS norm, aktiveerimisfunktsioonid jne). Seetõttu otsustasime nende rakendamiseks kasutada DSL-e: ehitasime ROSE esialgu Tritoni abil. Kuid sellest ajast alates oleme võtnud CuTeDSL-i kasutusele GPU programmeerimise peamise keskkonnana.

Matemaatilisest vaatenurgast on järeldustuuma tuumad tavaliselt lihtsad. Seetõttu ei seisne tuuma väljatöötamisel keerukus üldiselt mitte teostatava arvutuse väljendamises, vaid selles, kuidas riistvara peaks seda täpselt tegema. Parimate tuumade loomiseks vajame keelt, mis pakub õigeid abstraktsioone aluseks oleva riistvara kohta. Õige keel võimaldab meil määratleda arvutuse teostamiseks vajalikud kõrgetasemelised operatsioonid, spetsialiseerides ja optimeerides neid agressiivselt kõigis laia konfiguratsiooniruumi punktides.

CuTeDSL on Pythoni-põhine DSL NVIDIA GPU-de jaoks, mis on üles ehitatud CuTe paigutuse algebrale ja MLIR-ile. See kompileeritakse käigult (just-in-time) kõrgelt optimeeritud PTX-iks, säilitades samal ajal kontrolli madala taseme riistvara primitiivide üle. Selle eesmärk on pakkuda samu funktsioone ja jõudlust nagu CUTLASS C++, ilma kulukate kompileerimisaegadeta, mis pärsivad arenduskiirust. Kõigist keeltest, mida oleme uurinud või millega töötanud (sealhulgas CUDA, cuTile, TVM IR ja Triton), peame seda oma tuumadedevelopment-i jaoks kõige sobivamaks.

Üldiselt, kuigi see pole nii kõrgetasemeline ja kokkuvõtlik kui Triton või cuTile, leidsime, et CuTe paigutuse algebra ja selle kompositsioonid (mida pakutakse cute.Tensor kaudu) on mõistlikud asendajad. Lisaks pakub CuTeDSL piiramatut juurdepääsu riistvara primitiividele, tagades meile vajaliku kontrolli taseme tippjõudluse saavutamiseks kõigis äärmuslikes juhtudes (edge cases). CUDA puhul olid kompileerimisaegse spetsialiseerumise võimalused piiratud, kuna iga täiendav parameeter lisas kompileerimisaegadele ebamõistliku kulu või ehitussüsteemile keerukuse. Seevastu käigult kompileeritava (JIT) DSL-i abil saame spetsialiseeruda agressiivsemalt ilma iteratsioonitempot ohverdamata. Erinevalt TVM IR-ist, mis kasutas vaheastmena CUDA-d, leidsime, et PTX-i madalamale tasemele viiv MLIR-põhine torujuhe on tugevam ja hõlpsamini silutav kui see, mis väljastas vaheastmena CUDA-d. Lisaks leidsime kohandatud operatsioonide (custom ops) kasutamisega PTX-i juhiste kapseldamiseks, et suudame täielikult ära kasutada aluseks oleva riistvara kõiki funktsioone.

CUDA tuumade kompileerimiskuludest kulub tüüpiliselt suurem osa tuumade malli laiendamise faasile, kuna optimeerimine ja kompileerimine PTX etapist edasi on teostatav sekunditega. Paljude kompileerimisaegsete parameetrite korral võib see kompileerimise etapp mitte-tühise tuuma puhul paisuda mitmele minutile. Seda kitsaskohta võimendab asjaolu, et paralleelse kompileerimise võimalused tekivad alles hiljem, IR-i ja PTX-i tasemel, samas kui C++ mallide haldamine toimub järjestikuselt. Kuigi selle lahendamiseks on olemas nippe, mis seisnevad malli eksemplaride viimises eraldi kompileerimisüksustesse, nõuavad need olulisi muudatusi ehitussüsteemides ja kohmakate makrode kasutamist. Seevastu käigult (JIT) kompileeritavad CuTeDSL tuumad ei nõua keerulisi ehitusreegleid ega aeglusta kompileerimisaegu, kuna arendustsükkel saab keskenduda konfiguratsiooniruumi ühele punktile, samas kui hoolikad ühikutestid katavad hiljem tuuma kogu disainiruumi.

Ehkki JIT-kompileerimine nihutab kulud arendusajalt käitusaajale, on lisakulu enam kui aktsepteeritav. Mudeli juurutamisel aktiveerib see vaid mõned punktid kogu konfiguratsiooniruumist, tekitades järeldusmootori käivitamise/soojendamise ajal sekundi suurusjärgus lisakulu. Seevastu tuuma kõigi konfiguratsioonide eelkompileerimine võib kesta mitu minutit. Lisaks saab lisakulusid veelgi vähendada JIT-kompileeritud artefaktide vahemällu salvestamisega.

CUDA rasket kompileerimiskulu võimendab vigade päritolu jälgimise keerukus, kui kasutatakse mallipõhist suunamist ja spetsialiseerumist, nagu on omane CuTe-le ja CUTLASS-ile. DSL-i puhul saab vead suunata täpselt sellele koodireale, mis need tekitas, mis lihtsustab silumist märkimisväärselt. Lisaks on genereeritud koodi kaardistamine seda algatanud DSL-i lausele palju lihtsam vahepealsete MLIR-etappide kaudu või tuginedes genereeritud PTX-ile lisatud silumisteabele. Tänu nendele DSL-i aspektidele leidsime end tuumadearenduses palju produktiivsematena.

Kuigi tuumadel on mudelite vahel palju sarnasusi, peavad need siiski olema varustatud mitmesuguste konfiguratsiooninuppudega, mis võivad mõjutada jõudlust. Asjakohased parameetrid on järgmised:

Varjatud mõõde / Pea mõõde / Tähelepanupeade arv (Hidden Dimension/Head Dimension/Number of Attention Heads)

Sisend-/väljundandmetüübid

Kvantimisskeem

Kompileerimisaegsed konstandid (nihked, aktiveerimisfunktsioonid jne)

ROSE arhitektuur, sealhulgas selle mootor, mudel, tuumateek ja teenused.

Kompileerimisaegne spetsialiseerumine on kasulik nii jõudluse kui ka silumise jaoks. Eriti torujuhtmega protsessoris, nagu on omane GPU-dele, on olematu haru parem kui alati võetud või kunagi võtmata haru. Lisaks toob tuuma kompileerimine kindla konfiguratsiooninuppude komplekti jaoks kaasa vähem PTX/SASS-juhiseid, mida on uurimise käigus cuda-gdb kasutades lihtsam analüüsida ja siluda.

Antud pakendi suuruse ja konfiguratsiooni korral saab DRAM-i oodatavat liiklust ja aritmeetiliste operatsioonide arvu hinnata. Seejärel saab juhtumipõhiselt teha optimaalseid otsuseid optimaalse ruudustiku (grid) eraldamise, sünkroonimise ja teatud riistvarafunktsioonide kasutamise kohta. CuTeDSL võimaldab meil määratleda ja taaskasutada komponente sügavaks spetsialiseerumiseks ilma vastuvõetamatu kompileerimisaegse lisakuluta.

Üldiselt tuleb järeldustuumade jõudlust optimeerida kahe eri kasutusjuhu jaoks:

Eeltäide (Prefill): väike pakendi suurus (vähe tükke, jadasid), palju tokeneid jada kohta. Läbilaskvust tuleb maksimeerida nii manuste kui ka keelemudelite sisendtokenite puhul. Latentsus kümnetest sadade millisekunditeni.

Dekodeerimine (Decode): väikesest suurteni pakendisuurused, üks (ilma spekulatiivse dekodeerimiseta) kuni mõni token jada kohta (kavand-sihtmärk spekulatsioon, mitme tokeni ennustamine). Latentsus lähedane tuuma käivitamise latentsusele.

GEMM-i ja tähelepanu jaoks töötab eeltäite etapp suure hulga tokenitega ja jõudlus on piiratud seadme arvutusläbilaskvusega. Dekodeerimine hõlmab sama operatsiooni väikeste pakendisuurustega, kus täitmisaegu domineerib aktiveerimiste ja kaalude DRAM-ist voogprotsessoritesse ja tensorituumadesse laadimise kulu, mida piirab hoopis seadme mäluribaulaius. Ehkki pakendi suurus või tokenite üldarv pakendis on suhteliselt selge näitaja, on eeltäite ja dekodeerimise vaheline üleminekulävi peenem, kuna mäluribaulaiust määratlevad ka mudeli varjatud mõõde ja aktiveerimise andmetüüp. Tuumad jagatakse seega eeltäite- ja dekodeerimisrakendusteks, kui algoritmilised erinevused või ruudustiku/ploki määramised annaksid eelise neist kummalegi.

Parameeter, mis erineb universaalselt eeltäite- ja dekodeerimistuumade vahel, on lõimede arv plokis (lõimede arv jagatud lõimerühma (warp) suurusega 32 lõime lõimerühma kohta). Seda parameetrit on soovitav muuta kompileerimisaegseks konstandiks, kuna see võib aidata registri eraldamisel. CUDA-s tehakse seda tavaliselt lõimede arvu määramisega __launch_bounds__ kaudu ja lõimerühmade arvu esitamisega malliparameetrina:

cpp

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

}

Registrite arv, mis on saadaval ploki kõigile lõimedele, on riistvaraline piirang: 64K Hopperil ja Blackwellil. Kui lõimede arv on ette teada, teab kompilaator täpselt, mitu registrit ta saab igale lõimele määrata, ja suudab agressiivsemalt optimeeritud tuumad edukalt kompileerida. Sellega kaasneb siiski kompileerimisaegne kulu, kuna nüüd tuleb tuum erinevate lõimerühma suuruste jaoks ümber kompileerida. CuTeDSL-i puhul see enam probleemiks ei ole, kuna me ei pea maksma kogu konfiguratsiooniruumi kompileerimiskulu ette. Samuti eeldab DSL, et ploki suurus on konstantne avaldis, teostades kaudselt optimeerimise, mida CUDA-s eeldab __latch_bounds__.

Lisaks lõimerühmade arvu konstandiks muutmisele on soovitav seda eeltäite ja dekodeerimise vahel varieerida. GPU mäluribaulaiuse täielikuks ärakasutamiseks ei piisa alati suure lõimerühmade arvu määramisest ploki kohta, kuna täielik ribaulaius on saavutatav vaid võimalikult paljude saadaolevate voogprotsessorite (Streaming Multiprocessors, SM) kasutamisega. Seega valime antud tokenite arvu ning teadaolevate andmetüüpide ja varjatud mõõtmete korral esmalt ideaalse lõimepõhise vektori mõõtme (8 bfloat16 jaoks, 4 float32 jaoks), tagamaks, et genereeritud kood kasutab 128-bitiseid laadimis-/salvestusjuhiseid (256-bitiseid Blackwellil). Teades vektori suurust, saame arvutada tuuma jaoks vajalike lõimerühmade arvu. Püüame esmalt maksimeerida SM-ide arvu: kui lõimerühmi on vähem kui SM-e, valitakse ploki kohta üksainus lõimerühm. Tavaliselt on see tuuma dekodeerimisvariant. Kui lõimerühmi on piisavalt, suurendame lõimerühmade arvu järk-järgult 8, 16 või 32-ni, tehes võrdlusaluse ülempiiri leidmiseks, millest alates me tavaliselt enam jõudluse paranemist ei näe.

Kuigi Triton tegi mõned nendest optimeerimistest kapoti all, leidsime, et saavutame sarnaseid tulemusi CuTe paigutuse algebra hoolika kasutamisega, ilma et meid piiraksid kõrgetasemelise kompilaatori puudujäägid. Eriti alamlõimerühma horisontaalsete redutseerimiste korral leidsime, et CuTeDSL pakkus meile peeneteralist kontrolli, mis on vajalik optimaalse koodi genereerimiseks.

Teatud olukordades saab töö jaotamiseks võimalikult paljude plokkide vahel kasutada käsku cooperative_groups::this_grid().sync() kõigi osalevate plokkide sünkroonimiseks. Kehtivad piirangud: kuna plokki ei saa pärast täitmise alustamist SM-ist ajakavast eemaldada, ei tohi edasimineku tagamiseks plokkide arv ületada seadmel saadaolevate SM-ide arvu, tagamaks, et igal hetkel saavad kõik plokid sünkroonimises osalevaid globaalse mälu asukohti värskendada ja pärida. Seda piirangut jõustab cudaLaunchCooperativeKernel. Ruudustikutõkke kulu on Hopperil umbes 3 us, millele lisandub teatud olukordades kõigi varasemate mäluoperatsioonide läbipuhumise kulu kaudse mälutõkke kaudu. Üldiselt on see võrreldav tuuma kaheks käivitamiseks jagamise lisakuluga, mis viib jõudlustundliku disainiotsuseni. Vaatamata asjaolule, et programmiline sõltuv käivitamine (Programmatic Dependent Launch, PDL) võib käivituskulusid alates Hopperist vähendada, on dekodeerimistuumade puhul lisakulu vältimine seda väärt.

Kui üks töökoormustest nõuab rohkem plokke kui SM-ide arv, tuleks tuum jagada, tuginedes tuumade käivitamise kaudsele tõkkele. See kehtib peamiselt eeltäite kohta. Kui kõigi töökoormuste jaoks ideaalne plokkide arv mahub ruudustikku, eelistatakse ruudustikutõket, nagu see on omane väiksema tokenite arvuga töötavatele dekodeerimistuumadele.

CuTeDSL-is saab neid otsuseid teha hostipoolses käivituskoodis ning tuumasid saab vajadusel koostada ja ruudustikutõketega vaheldumisi paigutada, nagu allolevas näites illustreeritud. Sellise hübriidstrateegiaga saab dekodeerimistuumad optimeerida minimaalse latentsuse saavutamiseks ja eeltäite tuumad maksimaalse läbilaskvuse jaoks. Erinevalt CUDA-põhistest lahendustest saab seda loogikat määratleda kokkuvõtlikumalt ja väiksema kompileerimisaegse lisakuluga.

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

Meie MoE kõigi-kõigiga (all-to-all) tuumade puhul, nagu allpool üksikasjalikult kirjeldatud, aitas paindlikkus ehitada tuumasid ruudustikutõketega või ilma ning optimeerida neid kompileerimise ajal saavutada CuTeDSL-i kaudu märkimisväärset jõudluse kasvu.

Esimene tuum, mille CuTeDSL-is ehitasime, oli spetsiaalne RMS-normi tuum, mis on optimeeritud väikeste mõõtmete jaoks. Hiljutiste Qweni mudelite ise-tähelepanu (self-attention) kihis kasutatavat QK-normi rakendatakse peapõhiselt 64, 128 või 256 elemendile. Erinevalt tavalisest RMS-normist, mis töötab tuhandeid ületavate varjatud mõõtmetega, nõudes horisontaalset redutseerimist nii jagatud mälu kui ka registrite kaudu, vajab QK-norm redutseerimist ainult lõimerühmade vahel, kusjuures mitu pead on võimalik suruda ühteainsasse lõimerühma. Sellest sai kasutusjuht, kus saime kompileerimise ajal agressiivselt optimeerida, et leida parimad tokeni-lõimerühma määramised kõigi peamõõtmete ja andmetüüpide jaoks.

RMS-norm on üsna lihtne operatsioon, kuigi vektori ruutnormi arvutamiseks on vaja horisontaalset summat:

cpp

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

Kuna tähelepanupead on tavaliselt väikesed ja lõimepõhised operatsioonid on tugevalt vektoriseeritud, saab ruutnormi arvutada üksainus lõimerühm. Mäluribaulaiuse maksimeerimiseks peaks iga lõim töötlema 128 bitti sisendit, mis vastab 8 bfloat16 väärtusele. Seega on pea töötlemiseks piisav pool lõimerühmast ning ruutnormi saab määrata shuffle_sync_bfly abil teostatava horisontaalse redutseerimise teel. Alustades pea mõõtmest ja andmetüübist, saame määrata peade arvu lõimerühma kohta ning määrata argumendid ja redutseerimisvoorude arvu kompileerimise ajal.

Seejärel saame ehitada spetsiaalse tuuma suhteliselt väheste koodiridadega, tuginedes automaatselt vektoriseeritud laadimistele/salvestustele ja kasutades lõimerühma redutseerimiseks madalama taseme arhitektuurispetsiifilisi operatsioone. Tuum määrab, millise peaga (Q või K) töötada, ja tuvastab rakendatavad vastavad kaalud. See laadib peast terve vektori ja summeerib selle esmalt lokaalselt. Seejärel arvutatakse ruutnorm pea kohta olevate lõimede arvu põhjal. Lõpuks määratakse automaatselt vektoriseeritud RMS-norm ja kirjutatakse mällu tagasi. Mudelispetsiifilised parameetrid, nagu WEIGHT_BIAS ja EPS, volditakse konstantidena (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õrreldes tavalise RMS-normi tuumaga tagab see spetsiaalne variant 2–3 korda parema läbilaskvuse ja latentsuse nii eeltäite kui ka dekodeerimise puhul.

MoE saatmine ja kombineerimine segavad sisendtokenid ekspertidele, kellele need on suunatud, teostades kaalutud akumulatsiooni pärast seda, kui MoE GEMM-tuumad on teostanud asjakohased maatrikskorrutused. Jõudlus pakub erilist huvi, kuna erinevad tooted tuginevad erineva suurusega MoE mudelitele. Juurutustüüpide põhjal eristame optimeerimiseks erinevaid tuumasid:

TP=1: Kohalik tuum ühele GPU-le. Saatmine segab tokenid asjakohastele ekspertidele ja kombineerimine akumuleerib need. Kasutatakse peamiselt Qwen3 30B juurutuste jaoks H100-l.

DP=1, TP ≠ 1: Tensor-paralleelne, ühe sõlme tuum. Sisend on dubleeritud, saatmine valib praeguse astme (rank) tokenid ja kombineerimine akumuleerib need kõigis astmetes. Akumulatsioon NVLink-i või NVLS-i kaudu. Kasutatakse suurte eeltäitjate jaoks.

DP ≠ 1, NVLink: Andmeparalleelne tuum valikulise tensor-paralleelsusega. Iga aste saadab tokeneid igale teisele astmele, mis hostib spetsiifilisi ekspekte. Suunamisteave agregeeritakse, tokenid edastatakse NVIDIA NVLink-i kaudu. Kasutatakse andmeparalleelsete H200 või GB200 juurutuste jaoks.

DP ≠ 1, InfiniBand: Punkt-punkti (Peer-to-Peer) andmeparalleelne ja tensor-paralleelne tuum, mis töötab nii NVLink-i kui ka NVIDIA Quantum InfiniBandi kaudu. Kasutatakse mitme sõlmega Hopperi juurutuste jaoks.

Kohalik tuum loendab suunamistabeli alusel igale eksperdile määratud tokenite arvu ja määrab igale tokenile unikaalse indeksi igale eksperdile kuuluvate tokenite loendis. Pärast suunamise sammu kopeeritakse tokenid valikuliselt polsterdatud puhvritesse, et need edastada MoE GEMM-tuumale. Tokenite asukohti kasutades teostab kombineerimiskuum registrites akumulatsiooni, kirjutades lõplikud tokenid väljundpuhvrisse.

Saatmise faas erineb eeltäite ja dekodeerimise vahel: väheste tokenite puhul on lugemiseks vähe suunamisteavet, seega töötleme seda ühes plokis ja anname märku ülejäänud ruudustikust, kusjuures üksikud plokid kopeerivad tokenite tükke. Tokeni määramine määratakse jagatud mälus loendurite aatomilise suurendamise teel. Kuna eeltäide peab töötlema rohkem suunamisteavet, jagatakse suunamine eraldi tuumas mitme ploki vahel. Indeksid määratakse esmalt lokaalselt jagatud mälus olevate loendurite abil ja seejärel agregeeritakse kõigis globaalse mälu plokkides, et määrata iga tokeni positsioon. Selle tuuma ruudustik kinnitatakse (clamped) SM-ide arvule, et saaks kasutada ruudustikupõhiseid tõkkeid. Järgnev tuum teostab kopeerimised, käivitades selle palju suurema plokkide ja lõimerühmade arvu kaudu, et kasutada saadaolevat mäluribaulaiust täielikult ära.

Kombineerimiskuum on suhteliselt lihtne, kuna see teostab vektoriseeritud akumulatsiooni, tuginedes suunamisteabele ja saatmise arvutatud nihetele. CuTeDSL osutus siin väärtuslikuks, kuna arvutust sai väljendada lihtsate CuTe tensorite kaudu, mis kompileeruvad tõhusaks koodiks kõigi andmetüüpide ja varjatud mõõtmete jaoks.

Tensor-paralleelses juurutuses on suunamisteave ja sisendtokenid dubleeritud. Kohaliku saatmise tuuma saab täielikult taaskasutada väikese muudatusega, et suunata ja valida ainult need tokenid, mis on määratud praegusel astmel asuvatele ekspertidele. Mõningate parameetrite ja kompileerimisaegse spetsialiseerumisega taaskasutatakse kohaliku saatmise tuumasid.

Kuna astmete arv on tavaliselt väiksem kui ekspertide arv, kellele token suunatakse, võib keskmisel juhul eeldada, et iga tokeni vähemalt üks koopiatest suunatakse igale astmele. Selles olukorras saab kohaliku kombineerimise tuuma taaskasutada tokeni kohalike koopiate agregeerimiseks, millele järgneb kõigi-taandamine (all-reduce) summeerimiseks kõigis osalevates astmetes. Eeltäite tuuma puhul käivitame eraldi tuumad, kuid dekodeerimine sulatab mõlemad operatsioonid üheks käivituseks.

Kõigi-taandamise (all-reduce) jaoks kasutame tavalist NVLink-i, kui osaleb ainult kaks astet, lülitudes NVLS-ile, kui kaasatud on 4 või emm-kummast rohkem astet, et saavutada 10-20% läbilaskvuse kasv. Tänu DSL-ile saame jagada suurema osa taandamisloogikast ja partnerpuhvri vahetustest MoE ja tavaliste kõigi-taandamise tuumade vahel. See on olukord, kus riistvara lähedane kontroll on kasulik, kuna NVLS-i jaoks on vajalikud DSL-operatsioonid multimem.ld_reduce ja multimem.st üle. Taandamise enda rakendamine järgib suures osas NCCL-i testide all_reduce näidet.

Andmeparalleelsed tuumad tuginevad NVLink-ile tokenite agregeerimiseks mitme seadme jaoks. Saatmisfaasis suunavad partnerid enne tokenite vahetamist esmalt lokaalselt, et tuvastada iga saadetava tokeni sihtnihe igal astmel. Seejärel alustavad nad tokenite kopeerimist vastuvõtupuhvritesse, mis on eraldatud igale saatjale konfliktide vältimiseks. Valmimise signaliseerimiseks kasutatakse tokenipõhiseid vabastamis-salvestamise (release-store) kirjutusi. Vastuvõtja poolel suunamisteave agregeeritakse, et suunata kõik tokenid MoE GEMM-i jaoks ühte puhvrisse. Vastuvõtjad pärivad partneri kirjutatud lippu omandamis-laadimisega (acquire-load), et oodata kirjutuste lõpuleviimist.

Kombineerimise poolel kopeeritakse tokenid NVLink-i kaudu lähteastmetele, kusjuures tokenipõhine lipp näitab kättesaamist. Kombineerimiskuum on sarnane kohaliku ja ühe sõlme variantidega, väikeste erinevustega adresseerimises ja mõningase täiendava sünkroonimisloogikaga.

Selle tuuma rakendamiseks ehitasime PTX-kapselduste teegi, mis toob esile süsteemi- ja seadmetaseme ulatusega vabastamis-omandamise (release-acquire) laadimised ja salvestused. Parandades sünkroonimislippude granulaarsust ja eraldades rohkem plokke tokenite saatmiseks ja vastuvõtmiseks NVLink-i kaudu, saavutasime varasemate CUDA tuumadega võrreldes kahekordse paranemise.

Meie punkt-punkti tuumad saavutasid varem tipptasemel jõudluse Hopperi GPU-del, mis olid ühendatud ConnectX-7 InfiniBandi adapteritega. Oleme sellest ajast alates porditud kombineerimistuumad CuTeDSL-i edasise jõudluse kasvu saavutamiseks. Kuna CUDA versioonid tuginesid juba suurel määral kompileerimisaegsele spetsialiseerumisele, osutusid need meie repositooriumis kompileerimisaegseteks kitsaskohtadeks, sundides meid mõningaid parameetreid ära jätma. JIT-kompileeritud DSL-ile migreerumisega spetsialiseerusime rohkemate parameetrite jaoks ja timmisime ruudustiku eraldisi, tuues sisse täiendavaid erinevusi eeltäite ja dekodeerimise vahel ligikaudu 3% latentsuse kasvu nimel.

ROSE tuumade vaatenurgast jaguneb Blackwelli tugi kahte kategooriasse: arhitektuurist sõltumatud ja arhitektuurispetsiifilised tuumad. Arhitektuurispetsiifilised tuumad, mis tuginevad tensorituumadele, nõuavad kohanemiseks selgesõnalist tööd. Kuid tänu DSL-i kasutamisele kohanduvad arhitektuurist sõltumatud tuumad triviaalselt.

Blackwelli GPU-l on kuni 160 SM-i, erinevalt Hopperi 132-st. Enamiku tuumade jaoks on see kompileerimisaegne konstant ja hostipoolse käivituskoodi loogika võtab selle arvesse käivitatava optimaalse (eeltäide vs. dekodeerimine) tuuma määramisel. Lisaks tutvustab Blackwell ka 256-bitiseid laadimisjuhiseid, mida saab hõlpsasti väljastada vektori laiuse (CuTe tensorite viimase mõõtme) suurendamisega, et parandada mäluribaulaiuse ärakasutamist.

Blackwell toob sisse riistvaralise toe märkimisväärsele hulgale kvantimisskeemidele (MXFP8, MXFP4, NVFP4). Tugi nende jaoks rakendati peamiselt olemasolevate tuumade timmimisega FP8 plokipõhise skaleerimise jaoks. JIT-tugi on siin ülioluline, kuna uusi režiime saab toetada ilma kompileerimisaegset konfiguratsiooniruumi plahvatuslikult suurendamata. Varem leidsime, et Triton on ebamugav horisontaalsete redutseerimiste väljendamisel väheste elementide vahel, üle poole või veerandi lõimerühma. Peeneteralisem kontroll redutseerimiste üle oli ülioluline optimaalse jõudluse saavutamiseks NVFP4 ja vastavalt MXFP8/MXFP4 jaoks 16 või 32 elemendi absoluutse maksimumväärtuse arvutamisel.

Suhteliselt uue keelena puutusid keelemudelid oma eelkoolituse ajal CuTeDSL-i koodiga vähe kokku. Lisaks erines meie praeguse koodi stiil, kus tuumad on optimeeritud aktiveerimiste ja normide jaoks, veidi olemasolevatest CUTLASS-i või Quacki näidetest, mis olid tugevalt optimeeritud TMA ja arhitektuurispetsiifiliste maatrikskorrutusjuhiste jaoks. Punktini jõudmine, kus keelemudelid saavad tuumade kirjutamisel abiks olla, oli järkjärguline protsess.

Ehitasime käsitsi esialgse tuumade ja utiliidifunktsioonide korpuse, järgides järjepidevaid mustreid, enne kui saime tuumade teisendamisel agentidele tugineda. Korpuse arenedes saime teostada ligikaudseid tõlkeid CUDA-st CuTeDSL-i, millele järgnes märkimisväärne käsitsi ümbertöötamine. Eriti leidsime, et mudelid maadlevad osuti-tensoriks (pointer-to-tensor) teisenduste ja vektoriseerimisega, jäädes jänni tüügivigade parandamisega testide kompileerimisel. Nad on näidanud vähest arusaamist tuumade ruudustikupõhisest olemusest, ajades segamini lõime ja ruudustiku sünkroonimise. Kõige masendavam oli kiusatus hallutsineerida CUDA-spetsiifiliste teekide mittekasutuses olevaid tõlkeid meie endi CuTeDSL-i asenduste ehitamise asemel. Aja jooksul, kui CuTeDSL-i tuumade korpus kasvas, leidsime, et keelemudelid suudavad üles ehitada suurema osa sõustikukoodist (skeleton code) ja täita lüngad, võimaldades meil keskenduda kõrgema taseme disainile.

Võrreldes CUDA-ga leiame, et CuTeDSL pakub palju paremat kompileerimisjõudlust, võimaldades meil spetsialiseeruda rohkemate parameetrite jaoks ja saavutada jõudluse paranemist. Kuigi tuumad on Tritoni teatud abstraktsioonide puudumise tõttu veidi sõnaohtramad, leiame, et keerukus on õigustatud, kuna riistvara peeneteralisem kontroll võimaldab meil ehitada kiiremaid tuumasid.

Kirjutamise ajal on jõudluse vahe CUTLASS C++ ja CuTeDSL tuumade vahel minimaalne, avades põnevaid võimalusi suure jõudlusega GEMM-tuumade koostamiseks ja sulatamiseks seadmetevahelise side, normaliseerimise või aktiveerimistega. Pühendudes sellele DSL-ile, saame oma seadmetest välja pigistada täiendavat jõudlust, et parandada läbilaskvust kõigis meie töövoogudes, alates manustustest kuni keelemudeliteni.

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

Järeldustöötlus Perplexity AI-s

Miks CuTeDSL?

Arendajakogemus

Suur konfiguratsiooniruum

Eeltäida ja dekodeeri spetsialiseerumine

Ruudustiku sünkroonimine

CuTeDSL-i tuumad

QK norm

Ekspertide segu (Mixture-of-Experts) saatmine/kombineerimine (Dispatch/Combine)

Kohalik (Local)

Ühe sõlme (Singlenode)

Andmeparalleelne (Data Parallel)

Punkt-punkti (Peer-to-Peer)

Blackwelli tugi

AI-põhine tuumade genereerimine

Kokkuvõte ja tulevane töö

Viited

Viide