CuTeDSL uzņēmumā Perplexity

Uzņēmumā Perplexity mēs izmantojam savu iekšējo izpildlaika optimizēto apkalpošanas dzinēju (ROSE), lai apkalpotu modeļus, sākot no iegulījumiem līdz pat triljonu parametru LLM. Mūsu izpildlaiks, kas atrodas aiz visām API, tostarp Sonar, Search un Embeddings, tika izveidots, lai dotu iespēju pētniekiem izpētīt vispiemērotākos modeļus un arhitektūras, vienlaikus ātri izvietojot labākos kandidātus ražošanā. Lai to sasniegtu, mums ir jāspēj ātrā tempā ieviest, izvietot un optimizēt secinājumu iegūšanu jaunākajām arhitektūrām. Šajā emuāra ierakstā mēs pētām, kā CuTeDSL kodoli, kas integrēti mūsu secinājumu dzinējā, ļauj mums izveidot nepieciešamos GPU kodolus, lai nodrošinātu modeļu maksimālo veiktspēju uz NVIDIA GPU.

Pielāgoti modeļi, ko izveidojušas dažādas Perplexity komandas, tiek mitināti uzņēmumā NVIDIA Hopper un Blackwell GPU, izmantojot ROSE secinājumu dzinēju, ko izveidojusi AI secinājumu komanda. ROSE sākotnēji tika izveidots, lai labāk apkalpotu pielāgotus Llama modeļus gan valodas dekodēšanai, gan klasifikācijai, ar interfeisu, kas ir saderīgs ar NVIDIA Triton Inference Server. Laika gaitā tas izauga par pilnvērtīgu secinājumu dzinēju, kas apkalpo ievērojamus LLM, kā arī ranžēšanas, klasifikācijas, vērtēšanas un citus uz transformatoriem balstītus modeļus visdažādākajām komandām, lai nodrošinātu Perplexity produktu darbību. ROSE dizaina mērķis ir vienkāršība: laika gaitā tiek ieguldīts darbs, lai sasniegtu maksimālo veiktspēju, taču galvenā prioritāte ir pielāgoties pastāvīgi mainīgajām mūsu modeļu arhitektūrām un lietošanas gadījumiem.

Kodola konfigurācijas telpa pāri datu tipiem, galvu dimensijām, kompilēšanas laika konstantēm un palaišanas iestatījumiem.

Augstākajā līmenī ROSE sastāv no dzinēja, kas pielāgo modeli saskarnei, ar kuru klienti var mijiedarboties. Dzinēji ir ļoti atkārtoti izmantojami, tie apstrādā pieprasījumu plānošanu un pakešu apstrādi, kā arī izgūst dekodētos marķierus, iegulījumus vai rezultātus no pamatā esošajiem modeļiem. Tie nodrošina ierīču inicializāciju, starpaktuierīču un starpmezglu sakaru kanālus, kā arī svara ielādi. LLM gadījumā dzinējs ir atbildīgs par pakešu apstrādi, sagrupēšanu un paraugu ņemšanu. Tas ir atbildīgs arī par KV krātuves sadali gan pilnīgajai, gan lineārajai uzmanībai, kā arī par prefiksu saskaņošanu. Iegulījumu dzinēji ir vienkāršāki, jo tie veic tikai tiešsaistes pakešu apstrādi. Izpilde ir optimizēta tā, ka dzinējs sagatavo nākamo darba kārtu, kamēr modelis paralēli darbojas uz paātrinātāja.

ROSE piedāvā bagātīgu pielāgotu slāņu kolekciju modelēšanai, ietverot GPU kodolus, kas izveidoti ar CuTeDSL, Triton, CUDA, CUTLASS un cuBLAS. Matricu reizināšanai un uzmanības ieviešanai mēs galvenokārt paļaujamies uz NVIDIA kodoliem no CUTLASS un cuBLAS. Pārējā konveijera daļa, tostarp iegulījumu, normas, MoE un aktivizācijas kodoli, ir daļa no ROSE kodolu bibliotēkas. Lai vislabāk atbalstītu aparatūru, uz kuras mēs veicam secinājumu iegūšanu, ROSE ir aprīkots ar saviem ierīču starpsaziņas primitīviem MoE maršrutēšanai, nosūtīšanai un kombinēšanai. Lai gan citi kodoli ir salīdzinoši vienkārši, tiem jābūt rūpīgi specializētiem, lai sasniegtu maksimālu veiktspēju visās dažādajās slēptajās dimensijās un ierīcēs, uz kurām tie darbojas. Turklāt mēs paļaujamies uz kompilēšanas laika specializāciju, lai pielāgotos dažādu modeļu regulēšanas iespējām un pielāgojumiem (svara nobīde pret bez svara nobīdes RMS normai, aktivizācijas funkcijas utt.). Līdz ar to mēs nolēmām izmantot DSL, lai tos ieviestu: mēs sākotnēji izveidojām ROSE, izmantojot Triton. Tomēr kopš tā laika mēs esam pieņēmuši CuTeDSL kā galveno vidi GPU programmēšanai.

No matemātiskā viedokļa secinājumu kodoli parasti ir vienkārši. Līdz ar to, izstrādājot kodolu, sarežģītība parasti slēpjas nevis tajā, kā izteikt veicamo aprēķinu, bet gan tajā, tieši kā aparatūrai tas būtu jādara. Lai izveidotu labākos kodolus, mums ir nepieciešama valoda, kas piedāvā pareizās abstrakcijas virs pamatā esošās aparatūras. Pareizā valoda ļauj mums norādīt augsta līmeņa operācijas aprēķina veikšanai, agresīvi tās specializējot un optimizējot visos punktos plašā konfigurācijas telpā.

CuTeDSL ir uz Python balstīts DSL NVIDIA GPU, kas izveidots uz CuTe izkārtojuma algebras un MLIR bāzes. Tas kompilējas tieši laikā (just-in-time) uz augsti optimizētu PTX, vienlaikus saglabājot kontroli pār zema līmeņa aparatūras primitīviem. Tā mērķis ir piedāvāt tādas pašas funkcijas un veiktspēju kā CUTLASS C++, izvairoties no dārgajiem kompilēšanas laikiem, kas kavē izstrādes ātrumu. No valodām, kuras esam izpētījuši vai ar kurām esam strādājuši (tostarp CUDA, cuTile, TVM IR un Triton), mēs to uzskatām par vispiemērotāko mūsu kodolu izstrādei.

Kopumā, lai gan tas nav tik augsta līmeņa un lakonisks kā Triton vai cuTile, mēs atklājām, ka CuTe izkārtojuma algebra un tās kompozīcijas (kas eksponētas, izmantojot cute.Tensor) bija saprātīgi aizvietotāji. Turklāt CuTeDSL piedāvā neierobežotu piekļuvi aparatūras primitīviem, nodrošinot mums tādu kontroles līmeni, kāds nepieciešams, lai sasniegtu maksimālu veiktspēju visos malējos gadījumos. Ar CUDA kompilēšanas laika specializācijas iespējas bija ierobežotas, jo katrs papildu parametrs radīja nepamatotas izmaksas kompilēšanas laikam vai sarežģītību būvniecības sistēmai. Turpretim ar Just-In-Time kompilētu DSL mēs varam specializēties agresīvāk, neupurējot iterācijas tempu. Atšķirībā no TVM IR, kurš izmantoja CUDA kā starpposmu, mēs atklājām, ka uz MLIR balstītais konveijers, kas tiek samazināts līdz PTX, ir izturīgāks un vieglāk atkļūdojams nekā tas, kas emitēja CUDA kā starpposmu. Turklāt, izmantojot pielāgotas operācijas, lai aptvertu PTX instrukcijas, mēs atklājām, ka varam pilnībā izmantot visas pamatā esošās aparatūras funkcijas.

Lielākā daļa CUDA kodolu kompilēšanas izmaksu parasti tiek tērēta kodolu veidņu izvēršanas fāzē, jo optimizāciju un kompilāciju no PTX posma un tālāk var pabeigt dažu sekundju laikā. Ar daudziem kompilēšanas laika parametriem šis kompilēšanas posms var izplūst vairākās minūtēs nebūtiskam kodolam. Šo vājo vietu pastiprina fakts, ka paralēlas kompilācijas iespējas rodas tikai vēlāk, IR un PTX līmenī, savukārt C++ veidņu apstrāde notiek secīgi. Lai gan pastāv triki, lai to risinātu, pārvietojot veidņu instanciācijas uz atsevišķām kompilācijas vienībām, tiem ir nepieciešami būtiski pielāgojumi būvniecības sistēmās un neērtu makro izmantošana. Turpretim CuTeDSL kodoli, kas tiek kompilēti Just-In-Time (JIT), neprasa sarežģītus būvēšanas noteikumus un netraucē kompilēšanas laiku, jo izstrādes ciklu var koncentrēt uz vienu punktu konfigurācijas telpā, savukārt rūpīgi vienību testi vēlāk var aptver visu kodola dizaina telpu.

Pat ja JIT kompilācija pārvieto izmaksas no izstrādes laika uz izpildlaiku, pieskaitāmās izmaksas ir vairāk nekā pieņemamas. Kad modelis tiek izvietots, tas aktivizē tikai dažus punktus no visas konfigurācijas telpas, radot sekunžu mēroga pieskaitāmās izmaksas secinājumu dzinēja palaišanas/uzsilšanas laikā. Turpretim kodola visu konfigurāciju iepriekšēja kompilācija var ilgt vairākas minūtes. Turklāt pieskaitāmās izmaksas var vēl vairāk mazināt, saglabājot kešatmiņā JIT kompilētos artefaktus.

Smagās CUDA kompilēšanas izmaksas pastiprina grūtības izsekot kļūdu izcelsmei, ja tiek izmantota uz veidnēm balstīta nosūtīšana un specializācija, kā tas ir CuTe un CUTLASS gadījumā. Izmantojot DSL, kļūdas var norādīt uz koda līniju, kas to ģenerējusi, būtiski atvieglojot atkļūdošanu. Turklāt ir daudz vienkāršāk kartēt ģenerēto kodu uz DSL priekšrakumu, kas to radīja, izmantojot starpposma MLIR posmus vai paļaujoties uz atkļūdošanas informāciju, kas pievienota ģenerētajam PTX. Pateicoties šiem DSL aspektiem, mēs konstatējām, ka esam daudz produktīvāki kodolu izstrādē.

Lai gan kodoliem ir daudz līdzību starp modeļiem, tie joprojām ir jāaprīko ar virkni konfigurācijas pogu, kas var ietekmēt veiktspēju. Attiecīgie parametri ietver:

Slēptā dimensija / Galviņas dimensija / Uzmanības galviņu skaits

Ievades/izvades datu tipi

Kvantēšanas shēma

Kompilēšanas laika konstantes (nobīdes, aktivizācijas funkcijas utt.)

ROSE arhitektūra, tostarp tās dzinējs, modeļa, kodola bibliotēka un pakalpojumi.

Kompilēšanas laika specializācija ir noderīga gan veiktspējai, gan atkļūdošanai. Jo īpaši cauruļveida procesorā, kā tas ir GPU gadījumā, neeksistējošs zars ir labāks nekā vienmēr izmantots vai nekad neizmantots zars. Turklāt kodola kompilēšana konkrētai konfigurācijas pogu kopai rada mazāk PTX/SASS instrukciju, kuras ir vieglāk analizēt un atkļūdot, izmantojot cuda-gdb izmeklēšanas laikā.

Noteiktam pakešu lielumam un konfigurācijai var aprēķināt paredzamo trafiku uz DRAM un aritmētisko operāciju skaitu. Pēc tam optimālus lēmumus var pieņemt katrā gadījumā atsevišķi, lai nodrošinātu optimālu režģa sadali, sinhronizāciju un noteiktu aparatūras funkciju izmantošanu. CuTeDSL ļauj mums definēt un atkārtoti izmantot komponentus dziļai specializācijai, neradot nepieņemamas kompilēšanas laika pieskaitāmās izmaksas.

Parasti secinājumu kodolu veiktspēja ir jāoptimizē diviem dažādiem lietošanas gadījumiem:

Iepriekšēja uzpilde (Prefill): mazs pakešu lielums (maz gabalu, secību), daudz marķieru uz secību. Caurlaides spēja ir jāmaksimizē gan iegulījumiem, gan LLM ievades marķieriem. Aizkave desmitos līdz simtos milisekunžu.

Dekodēšana (Decode): mazi līdz lieli pakešu lielumi, viens (bez spekulatīvas dekodēšanas) līdz dažiem marķieriem uz secību (melnraksta-mērķa spekulācijas, vairāku marķieru prognozēšana). Aizkave tuvu kodola palaišanas aizkavei.

GEMM un uzmanībai iepriekšējas uzpildes posms darbojas ar lielu marķieru skaitu, un veiktspēju ierobežo ierīces skaitļošanas caurlaides spēja. Dekodēšana ietver to pašu operāciju ar maziem pakešu lielumiem, un izpildes laiku dominē aktivizāciju un svaru ielādes izmaksas no DRAM straumēšanas daudzprocesoros un tensora kodolos, tā vietā ierobežojot ar ierīces atmiņas joslas platumu. Lai gan pakešu lielums vai kopējais marķieru skaits paketē ir salīdzinoši skaidrs rādītājs, pārslēgšanās slieksnis starp iepriekšēju uzpildi un dekodēšanu ir smalkāks, jo atmiņas joslas platumu nosaka arī modeles slēptā dimensija un aktivizācijas datu tips. Līdz ar to kodoli tiek sadalīti iepriekšējas uzpildes un dekodēšanas ieviešanās, ja algoritmiskās atšķirības vai režģa/bloka uzdevumi dotu priekšroku kādam no tiem.

Parametrs, kas universāli atšķiras starp iepriekšējas uzpildes un dekodēšanas kodoliem, ir pavedienu grupu (warp) skaits blokā (pavedienu skaits, dalīts ar pavedienu grupas lielumu — 32 pavedieniem uz grupu). Šo parametru vēlams padarīt par kompilēšanas laika konstanti, jo tas var palīdzēt reģistru sadalē. CUDA tas parasti tiek darīts, norādot pavedienu skaitu, izmantojot __launch_bounds__, un nodrošinot pavedienu grupu skaitu kā veidnes parametru:

cpp

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

}

Reģistru skaits, kas pieejams visiem bloka pavedieniem, ir aparatūras ierobežojums — 64K uz Hopper un Blackwell. Ja pavedienu skaits ir zināms iepriekš, kompilators precīzi zina, cik reģistru tas var piešķirt katram pavedienam, un var veiksmīgi kompilēt agresīvāk optimizētus kodolus. Tomēr tas notiek par kompilēšanas laika izmaksām, jo tagad kodols ir jāpārakompilē dažādiem pavedienu grupu izmēriem. Ar CuTeDSL tā vairs nav problēma, jo mums nav jāmaksā kompilēšanas izmaksas par visu konfigurācijas telpu iepriekš. DSL arī paredz, ka bloka izmēram ir jābūt konstantai izteiksmei, netieši veicot optimizāciju, ko CUDA implicē __latch_bounds__.

Papildus tam, ka pavedienu grupu skaits tiek padarīts par konstanti, to vēlams mainīt starp iepriekšējo uzpildi un dekodēšanu. Lai pilnībā izmantotu GPU atmiņas joslas platumu, nepietiek vienmēr piešķirt lielu pavedienu grupu skaitu uz bloku, jo pilnu joslas platumu var sasniegt tikai, izmantojot pēc iespējas vairāk pieejamo straumēšanas daudzprocesoru (SM). Līdz ar to noteiktam marķieru skaitam un zināmiem datu tipiem un slēptajām dimensijām mēs vispirms izvēlamies ideālo vektora dimensiju uz pavedienu (8 priekš bfloat16, 4 priekš float32), lai nodrošinātu, ka ģenerētais kods izmanto 128 bitu ielādes/saglabāšanas instrukcijas (256 bitu uz Blackwell). Zinot vektora izmēru, mēs varam aprēķināt kodolam nepieciešamo pavedienu grupu skaitu. Mēs vispirms cenšamies maksimizēt SM skaitu: ja ir mazāk pavedienu grupu nekā SM, tiek izvēlēta viena pavediena grupa uz bloku. Parasti tā ir kodola dekodēšanas varianta. Ja ir pietiekami daudz pavedienu grupu, mēs pakāpeniski palielinām pavedienu grupu skaitu līdz 8, 16 vai 32, veicot salīdzinošo testēšanu līdz augšējai robežai, pēc kuras mēs parasti vairs neredzam veiktspējas uzlabojumu.

Lai gan Triton veica dažas no šīm optimizācijām automātiski, mēs atklājām, ka mēs varam sasniegt līdzīgus rezultātus, rūpīgi izmantojot CuTe izkārtojuma algebru, netiekot ierobežotiem ar augsta līmeņa kompilatora trūkumiem. Jo īpaši, ja ir iesaistītas apakšgrupu horizontālās redukcijas, mēs atklājām, ka CuTeDSL piedāvāja mums smalki graudainu kontroli, kas nepieciešama optimāla koda ģenerēšanai.

Dažās situācijās, lai izplatītu darbu pa pēc iespējas vairāk blokiem, var izmantot cooperative_groups::this_grid().sync(), lai sinhronizētu visus iesaistītos blokus. Tiek piemēroti ierobežojumi: tā kā bloku nevar atcelt no SM plānošanas, kad tas sāk izpildi, lai garantētu virzību uz priekšu, bloku skaits nedrīkst pārsniegt ierīcē pieejamo SM skaitu, lai nodrošinātu, ka jebkurā brīdī visi bloki var atjaunināt un aptaujāt globālās atmiņas vietas, kas iesaistītas sinhronizācijā. Šo ierobežojumu uzspiež cudaLaunchCooperativeKernel. Režģa barjeras izmaksas uz Hopper ir aptuveni 3us, pieskaitot visu iepriekšējo atmiņas operāciju skalošanas izmaksas, izmantojot netiešu atmiņas barjeru dažās situācijās. Kopumā tas ir salīdzināms ar kodola sadalīšanas divās palaišanās reizēs pieskaitāmajām izmaksām, kas noved pie veiktspējas jūtīga dizaina lēmuma. Neskatoties uz to, ka programmējamā atkarīgā palaišana (PDL) var samazināt palaišanas pieskaitāmās izmaksas no Hopper un uz priekšu, dekodēšanas kodoliem ir vērts izvairīties no šīm pieskaitāmajām izmaksām.

Ja viena no darba slodzēm prasa vairāk bloku nekā SM skaits, kodols ir jāsadala, paļaujoties uz kodola palaišanas netiešo barjeru. Tas galvenokārt attiecas uz iepriekšēju uzpildi. Ja bloku skaits, kas ir ideāls visām darba slodzēm, ietilpst režģī, priekšroka tiek dota režģa barjerai, kā tas ir dekodēšanas kodolu gadījumā, kas darbojas ar mazāku marķieru skaitu.

CuTeDSL šādus lēmumus var pieņemt resursdatora puses palaišanas kodā, un kodolus var kompozēt un mijiedarboties ar režģa barjerām pēc vajadzības, kā ilustrēts zemāk redzamajā piemērā. Izmantojot šādu hibrīda stratēģiju, dekodēšanas kodolus var optimizēt, lai sasniegtu minimālu aizkavi, un iepriekšējas uzpildes kodolus var optimizēt maksimālai caurlaides spējai. Atšķirībā no uz CUDA balstītiem risinājumiem, šo loģiku var definēt lakoniskāk un ar mazākām kompilēšanas laika pieskaitāmajām izmaksām.

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ūsu MoE all-to-all kodolu gadījumā, kā sīkāk aprakstīts zemāk, elastība izveidot kodolus ar vai bez režģa barjerām un optimizēt tos kompilēšanas laikā palīdzēja mums sasniegt ievērojamus veiktspējas ieguvumus, izmantojot CuTeDSL.

Pirmais kodols, ko izveidojām CuTeDSL, bija specializēts RMS normas kodols, kas optimizēts mazām dimensijām. QK norma, ko izmanto jaunāko Qwen modeļu pašas uzmanības slānī, tiek lietota uz galviņas pamata 64, 128 vai 256 elementiem. Atšķirībā no parastās RMS normas, kas darbojas ar slēptajām dimensijām, kuras pārsniedz tūkstošus, pieprasot horizontālu redukciju gan caur koplietojamo atmiņu, gan reģistriem, QK normai nepieciešama redukcija tikai starp pavedienu grupām, ar iespēju saspiest vairākas galviņas vienā pavedienu grupā. Tas kļuva par lietošanas gadījumu, kurā mēs varējām agresīvi optimizēt kompilēšanas laikā, lai atrastu labākos marķieru un pavedienu grupu sadalījumus visām galviņu dimensijām un datu tipiem.

RMS norma ir diezgan vienkārša operācija, lai gan tā prasa horizontālu summu, lai aprēķinātu vektora kvadrātisko normu:

cpp

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

Tā kā uzmanības galviņas parasti ir mazas un operācijas uz pavedienu ir stipri vektorizētas, viena pavediena grupa var aprēķināt kvadrātisko normu. Lai maksimizētu atmiņas joslas platumu, katram pavedienam vajadzētu apstrādāt 128 bitu ievadi, kas nozīmē 8 bfloat16 vērtības. Līdz ar to pusgrupa ir pietiekama, lai apstrādātu galviņu, un kvadrātisko normu var noteikt, izmantojot horizontālo redukciju, ko veic shuffle_sync_bfly. Sākot no galviņas dimensijas un datu tipa, mēs varam noteikt galviņu skaitu uz pavedienu grupu un kompilēšanas brīdī norādīt argumentus un redukcijas kārtu skaitu.

Mēs tad varam izveidot specializētu kodolu ar salīdzinoši nedaudzām koda rindiņām, paļaujoties uz automātiski vektorizētām ielādēm/saglabāšanām un izmantojot zemāka līmeņa arhitektūrai specifiskas operācijas pavedienu grupas redukcijai. Kodols nosaka, ar kuru galviņu (Q vai K) strādāt, un identificē atbilstošos svirus, ko lietot. Tas ielādē veselu vektoru no galviņas un vispirms summē lokāli. Pēc tam, pamatojoties uz pavedienu skaitu uz galviņu, tiek aprēķināta kvadrātiskā norma. Visbeidzot, RMS norma, kas ir automātiski vektorizēta, tiek noteikta un ierakstīta atpakaļ atmiņā. Modeļam specifiski parametri, piemēram, WEIGHT_BIAS un EPS, ir konstantes (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)

Salīdzinājumā ar parasto RMS normas kodolu, šis specializētais variants nodrošina 2-3 reizes labāku caurlaides spēju un aizkavi gan iepriekšējai uzpildei, gan dekodēšanai.

MoE nosūtīšana un kombinēšana samaisa ievades marķierus uz ekspertiem, uz kuriem tie ir maršrutēti, veicot svērtu akumulāciju pēc tam, kad MoE GEMM kodoli veic attiecīgās matricas reizināšanas. Veiktspēja ir īpaši svarīga, jo dažādi produkti balstās uz dažāda izmēra MoE modeļiem. Pamatojoties uz izvietošanas veidiem, mēs izšķiram dažādus optimizējamos kodolus:

TP=1: Lokālais kodols vienam GPU. Nosūtīšana samaisa marķērus uz attiecīgajiem ekspertiem, un kombinēšana tos uzkrāj. Galvenokārt izmanto Qwen3 30B izvietojumiem uz H100.

DP=1, TP ≠ 1: Tensora paralēlais, viena mezgla kodols. Ievade tiek replicēta, nosūtīšana atlasa marķierus pašreizējam rangam, un kombinēšana tos uzkrāj visos rangos. Uzkrāšana visā NVLink vai NVLS. Izmanto lieliem iepriekšējiem uzpildītājiem.

DP ≠ 1, NVLink: Datu paralēlais kodols ar izvēles tensora paralēlismu. Katrs rangs sūta marķierus katram citam rangam, kas mitina konkrētus ekspertus. Maršrutēšanas informācija tiek apkopota, marķieri tiek pārsūtīti pa NVIDIA NVLink. Izmanto datu paralēliem H200 vai GB200 izvietojumiem.

DP ≠ 1, InfiniBand: Peer-to-Peer datu paralēlais un tensora paralēlais kodols, kas darbojas gan starp NVLink, gan NVIDIA Quantum InfiniBand. Izmanto vairāku mezglu Hopper izvietojumiem.

Lokālais kodols saskaita marķieru skaitu, kas piešķirts katram ekspertam, pamatojoties uz maršrutēšanas tabulu, un piešķir unikālu indeksu katram marķierim to marķieru sarakstā, kas pieder katram ekspertam. Pēc maršrutēšanas solis marķieri tiek kopēti izvēles polsterētos buferos, lai tos nodotu tālāk MoE GEMM kodolam. Izmantojot marķieru atrašanās vietas, kombinēšanas kodols pēc tam veic akumulāciju reģistros, ierakstot galīgos marķierus izvades buferī.

Nosūtīšanas fāze atšķiras starp iepriekšējo uzpildi un dekodēšanu: dažiem marķieriem ir maz lasāmas maršrutēšanas informācijas, tāpēc mēs to apstrādājam vienā blokā un signalizējam par pārējo režģi, atsevišķiem blokiem kopējot marķieru gabalus. Marķiera piešķiršana tiek noteikta, atomāri palielinot skaitītājus koplietojamajā atmiņā. Tā kā iepriekšējai uzpildei ir jāapstrādā vairāk maršrutēšanas informācijas, maršrutēšana tiek sadalīta vairākos blokos atsevišķā kodolā. Indeksi vispirms tiek piešķirti lokāli, izmantojot skaitītājus koplietojamajā atmiņā, un pēc tam tiek apkopoti visos blokos globālajā atmiņā, lai noteiktu katra marķiera pozīciju. Šī kodola režģis ir fiksēts pie SM skaita, lai varētu izmantot režģa mēroga barjeras. Nākamais kodols veic kopijas, kas tiek palaistas daudz lielākā bloku un pavedienu grupu skaitā, lai pilnībā izmantotu pieejamo atmiņas joslas platumu.

Kombinēšanas kodols ir salīdzinoši vienkāršs, jo tas veic vektorizētu akumulāciju, paļaujoties uz maršrutēšanas informāciju un nobīdēm, ko aprēķinājusi nosūtīšana. CuTeDSL šeit izrādījās vērtīgs, jo aprēķinu varēja izteikt ar vienkāršiem CuTe tensoriem, kas kompilējas efektīvā kodā visiem datu tipiem un slēptajām dimensijām.

Tensora paralēlā izvietojumā maršrutēšanas informācija un ievades marķieri tiek replicēti. Lokālās nosūtīšanas kodolu var pilnībā izmantot atkārtoti ar nelielu izmaiņu, lai maršrutētu un atlasītu tikai tos marķierus, kas piešķirti ekspertiem, kuri atrodas pašreizējā rangā. Ar dažiem parametriem un kompilēšanas laika specializāciju lokālās nosūtīšanas kodoli tiek izmantoti atkārtoti.

Tā kā rangu skaits parasti ir mazāks par ekspertu skaitu, uz kuriem marķieris tiek maršrutēts, vidējā gadījumā var pieņemt, ka katram marķierim vismaz viena no tā kopijām tiek maršrutēta uz katru ranku. Šajā situācijā lokālās kombinēšanas kodols var tikt izmantots atkārtoti, lai apkopotu marķiera lokālās kopijas, kam seko all-reduce, lai summētu visus iesaistītos rangus. Iepriekšējas uzpildes kodola gadījumā mēs palaižam atsevišķus kodolus, tomēr dekodēšana apvieno abas operācijas vienā palaišanā.

All-reduce gadījumā mēs izmantojam vienkāršu NVLink, ja piedalās tikai divi rangi, pārslēdzoties uz NVLS, ja ir iesaistīti 4 vai vairāk rangi, lai nodrošinātu 10-20% caurlaides spējas pieaugumu. Pateicoties DSL, mēs varam kopīgot lielāko daļu redukcijas loģikas un vienaudžu bufera apmaiņas starp MoE un parastajiem all-reduce kodoliem. Šī ir situācija, kurā cieša kontrole pār aparatūru ir izdevīga, jo NVLS ir nepieciešamas DSL operācijas virs multimem.ld_reduce un multimem.st. Pašas redukcijas ieviešana lielā mērā seko all_reduce piemēram no NCCL testiem.

Datu paralēlie kodoli balstās uz NVLink, lai apkopotu marķierus vairākām ierīcēm. Nosūtīšanas fāzē, pirms marķieru apmaiņas, vienaudži vispirms maršrutē lokāli, lai identificētu mērķa nobīdi katrā rangā katram marķierim, ko tie sūta. Pēc tam tie sāk kopēt marķierus saņemšanas buferos, kas piešķirti katram sūtītājam, lai izvairītos no konfliktiem. Lai signalizētu par pabeigšanu, tiek izmantoti katra marķiera release-store ieraksti. Saņemšanas galā maršrutēšanas informācija tiek apkopota, lai varētu kondensēt visus marķierus vienā buferī MoE GEMM. Saņēmēji aptaujā vienaudžu ierakstīto karodziņu ar acquire-load, lai gaidītu rakstīšanas pabeigšanu.

Kombinēšanas pusē marķieri tiek kopēti pa NVLink uz avota rangiem, ar katra marķiera karodziņu, kas norāda saņemšanu. Kombinēšanas kodols ir līdzīgs lokālajiem un viena mezgla variantiem, ar nelielām atšķirībām adresēšanā un dažām papildu sinhronizācijas loģikām.

Mēs izveidojām PTX iesaiņotāju bibliotēku, lai eksponētu release-acquire ielādes un saglabāšanas operācijas gan sistēmas, gan ierīces līmenī, lai ieviestu šo kodolu. Uzlabojot sinhronizācijas karodziņu granularitāti un piešķirot vairāk bloku marķieru sūtīšanai un saņemšanai pa NVLink, mēs sasniedzām 2 reižu uzlabojumu salīdzinājumā ar mūsu iepriekšējiem CUDA kodoliem.

Mūsu Peer-to-Peer kodoli iepriekš sasniedza modernāko veiktspēju uz Hopper GPU, kas savienoti pārī ar ConnectX-7 InfiniBand adapteriem. Kopš tā laika esam portējuši kombinēšanas kodolus uz CuTeDSL, lai iegūtu papildu veiktspējas ieguvumus. Tā kā CUDA versijas jau smagi balstījās uz kompilēšanas laika specializāciju, tās izrādījās kompilēšanas laika vājās vietas mūsu repozitorijā, piespiežot mūs izlaist dažus parametrus. Migrējot uz JIT kompilēto DSL, mēs specializējāmies vairākiem parametriem un pielāgojām režģa sadales, ieviešot turpmākas atšķirības starp iepriekšējo uzpildi un dekodēšanu, lai nodrošinātu aptuveni 3% aizkaves uzlabojumu.

No ROSE kodolu viedokļa Blackwell atbalsts iedalās divās kategorijās: arhitektūrai agnostiskos un arhitektūrai specifiskos kodolos. Arhitektūrai specifiskie kodoli, kas paļaujas uz tensora kodoliem, prasa skaidru darbu pielāgošanai. Tomēr, pateicoties DSL izmantošanai, arhitektūrai agnostiskie kodoli pielāgojas banāli.

Blackwell GPU ir līdz 160 SM, salīdzinājumā ar Hopper 132. Vairumam kodolu tas ir kompilēšanas laika konstante, un loģika resursdatora puses palaidēja kodā to ņem vērā, nosakot optimālo (iepriekšējas uzpildes pret dekodēšanas) kodolu, ko palaist. Turklāt Blackwell ievieš arī 256 bitu ielādes instrukcijas, kuras var viegli emitēt, palielinot vektora platumu (CuTe tensoru pēdējo dimensiju), lai uzlabotu atmiņas joslas platuma izmantošanu.

Blackwell ievieš aparatūras atbalstu ievērojamam skaitam kvantēšanas shēmu (MXFP8, MXFP4, NVFP4). Atbalsts tām tika ieviests galvenokārt, pielāgojot esošos kodolus FP8 bloku mēroga mērogošanai. JIT atbalsts šeit ir izšķirošs, jo jaunus režīmus var atbalstīt, neizraisot kompilēšanas laika konfigurācijas telpas eksploziju. Agrāk mēs atklājām, ka Triton ir neveikls horizontālo redukciju izteikšanai pāri dažiem elementiem, virs puses vai ceturtdaļas pavedienu grupas. Smalkāka kontrole pār redukcijām bija izšķiroša, lai sasniegtu optimālu veiktspēju, aprēķinot absolūto maksimālo vērtību pāri 16 vai 32 elementiem attiecīgi NVFP4 un MXFP8/MXFP4.

Kā salīdzinoši jauna valoda, LLM pirmsapmācības laikā maz saskārās ar CuTeDSL kodu. Turklāt mūsu pašreizējā koda stils ar kodoliem, kas optimizēti aktivizācijām un normām, nedaudz atšķīrās no esošajiem CUTLASS vai Quack piemēriem, kas bija stipri optimizēti TMA un arhitektūrai specifiskām matricas reizināšanas instrukcijām. Nonākšana līdz punktam, kurā LLM varētu palīdzēt kodolu rakstīšanā, bija pakāpenisks process.

Mēs manuāli izveidojām sākotnējo kodolu un palīgfunkciju korpusu, sekojot konsekventiem modeļiem, pirms varējām paļauties uz aģentiem kodolu konvertēšanā. Attīstoties korpusam, mēs varējām veikt aptuvenus tulkojumus no CUDA uz CuTeDSL, kam sekoja ievērojams manuāls atkārtots darbs. Jo īpaši mēs atklājām, ka modeļi cīnījās ar rādītāju uz tensora pārvēršanu un vektorizāciju, iesprūstot, mēģinot labot tipa kļūdas, kad mērķis bija panākt testa kompilēšanos. Tie ir izrādījuši maz izpratnes par kodolu režģa raksturu, sajaucot pavedienu un režģa sinhronizāciju. Visnomācošāk bija kārdinājums halucinēt neesošus CUDA specifisko bibliotēku tulkojumus, tā vietā, lai veidotu pret mūsu pašu CuTeDSL aizvietotājiem. Laika gaitā, augot CuTeDSL kodolu korpusam, mēs atklājām, ka LLM var izveidot lielāko daļu koda skeleta un aizpildīt tukšumus, ļaujot mums koncentrēties uz augsta līmeņa dizainu.

Salīdzinājumā ar CUDA, mēs uzskatām, ka CuTeDSL nodrošina daudz labāku kompilēšanas laika veiktspēju, ļaujot mums specializēties vairākiem parametriem un sasniegt veiktspējas uzlabojumus. Lai gan kodoli ir nedaudz garāki dažu abstrakciju trūkuma dēļ salīdzinājumā ar Triton, mēs uzskatām, že sarežģītība ir pamatota, jo smalkāka kontrole pār aparatūru ļauj mums izveidot ātrākus kodolus.

Rakstīšanas laikā veiktspējas atšķirība starp CUTLASS C++ un CuTeDSL kodoliem ir minimāla, paverot aizraujošas iespējas kompozēt un sapludināt augstas veiktspējas GEMM kodolus ar starpaktuierīču saziņu, normalizāciju vai aktivizācijām. Apņemoties izmantot šo DSL, mēs varēsim iegūt papildu veiktspēju no mūsu ierīcēm, lai uzlabotu caurlaides spēju visās mūsu darba slodzēs, sākot no iegulījumiem līdz 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}
}

Secinājumu iegūšana uzņēmumā Perplexity AI

Kpēc CuTeDSL?

Izstrādātāju pieredze

Liela konfigurācijas telpa

Iepriekšēja aizpilde un dekodēšanas specializācija

Režģa sinhronizācija

CuTeDSL kodoli

QK norma

Mixture-of-Experts nosūtīšana/kombinēšana (Dispatch/Combine)

Lokālais (Local)

Viena mezgla (Singlenode)

Datu paralēls (Data Parallel)

Vienaudzis-pret-vienaudzi (Peer-to-Peer)

Blackwell atbalsts

Ar mākslīgo intelektu atbalstīta kodolu ģenerēšana

Secinājumi un turpmākais darbs

Atsauces

Citāts