CuTeDSL hjá Perplexity

Hvernig við notum CuTeDSL til að smíða GPU kjarna sem ná hámarksafköstum úr líkönum okkar á NVIDIA GPUskjákortum.

HöfundarPerplexity Research

Hjá Perplexity notum við okkar eigið Runtime-Optimized Serving Engine (ROSE) til að keyra líkön allt frá innfellingum yfir í trilljóna færibreytu LLM. Keyrsluumhverfið okkar, sem er á bak við öll forritaskil (API), þar á meðal Sonar, Search og Embeddings, var byggt til að gera rannsakendum kleift að kanna hentugustu líkönin og arkitektúrana, en jafnframt til að taka bestu frambjóðendurna hratt í notkun. Til að ná þessu þurfum við að geta útfært, sett upp og fínstillt ályktun (inference) fyrir nýjustu arkitektúra á örskotsstundu. Í þessari bloggfærslu skoðum við hvernig CuTeDSL kjarnar, sem eru samþættir í ályktunarvél okkar, gera okkur kleift að smíða þá GPU kjarna sem þarf til að ná hámarksafköstum úr líkönum á NVIDIA GPUskjákortum.

Ályktun hjá Perplexity AI

Sérsniðin líkön, sem eru smíðuð af ýmsum teymum hjá Perplexity, eru hýst innanhúss á NVIDIA Hopper og Blackwell GPUskjákortum með ROSE ályktunarvélinni sem AI Inference teymið smíðaði. ROSE var upphaflega smíðað til að þjónusta sérsniðin Llama-líkön betur fyrir bæði tungumálaafkóðun og flokkun, með viðmóti sem er samhæft við NVIDIA Triton Inference Server. Með tímanum þróaðist það í alvöru ályktunarvél sem þjónustar stór LLM, auk röðunar, flokkunar, einkunnagjafar og annarra spenni-undirstaða líkana (transformer-based models) fyrir fjölbreytt úrval teyma til að knýja vörur Perplexity. Hönnunarmarkmið ROSE er að vera einfalt: með tímanum er fjárfest í vinnu til að ná hámarksafköstum, en aðalforgangsverkefnið er að aðlagast stöðugt breytilegum arkitektúrum og notkunartilvikum fyrir líkönin okkar.

Vinnusvæði kjarnastillinga yfir gagnategundir, hausvíddir, samsetningarfastar og keyrslustillingar.

Á grófari stigum samanstendur ROSE af vél sem aðlagar líkan að viðmóti sem viðskiptavinir geta haft samskipti við. Vélarnar eru mjög endurnýtanlegar, sjá um beiðnaáætlun og lotuvinnslu (batching), og sækja afkóðuð tákn, innfellingar eða einkunnir frá undirliggjandi líkönum. Þær sjá um frumstillingu tækja, samskiptaleiða milli ferla og hnúta, og hleðslu á vægi. Fyrir LLM sér vélin um lotuvinnslu, skiptingu í hluta (chunking) og úrtak (sampling). Hún sér einnig um úthlutun KV-geymslu, bæði fyrir fulla og línulega athygli, auk forskeiddu samræmingar (prefix matching). Innfellingarvélar eru einfaldari, þar sem þær framkvæma aðeins lotuvinnslu í rauntíma. Keyrslan er fínstillt þannig að vélin undirbýr næstu umferð vinnu á meðan líkanið keyrir samhliða á hraðallinum (accelerator).

ROSE býður upp á mikið safn af sérsniðnum lögum fyrir líkanagerð, sem vefja utan um GPU kjarna byggða með CuTeDSL, Triton, CUDA, CUTLASS og cuBLAS. Fyrir fylkjamargföldun og útfærslur á athygli treystum við að mestu á NVIDIA kjarna frá CUTLASS og cuBLAS. Afgangurinn af pípunni (pipeline), þar á meðal innfellingar-, norm-, MoE- og virkjunarkjarnar (activation kernels), eru hluti af kjarnasafni ROSE. Til að styðja sem best við vélbúnaðinn sem við keyrum ályktun á er ROSE búið sínum eigin samskiptaeiningum á milli tækja fyrir MoE-beiningu, útsendingu (dispatch) og samsetningu (combine). Þótt hinir kjarnarnir séu hlutfallslega einfaldir verða þeir að vera vandlega sérhæfðir til að ná hámarksafköstum á öllum hinum ýmsu faldu víddum og tækjum sem þeir keyra á. Að auki treystum við á sérhæfingu á þýðingartíma til að styðja við allar stillingar og breytingar á ýmsum líkönum (vægishlutdrægni á móti engri vægishlutdrægni fyrir RMS norm, virkni föll o.s.frv.). Þess vegna ákváðum við að nota DSL til að útfæra þá: við byggðum ROSE upphaflega með Triton. Við höfum þó síðan tekið upp CuTeDSL sem aðalumhverfi fyrir GPU forritun.

Af hverju CuTeDSL?

Frá stærðfræðilegu sjónarmiði eru ályktunarkjarnar venjulega óflóknir. Þess vegna liggur flóknin við þróun kjarna almennt ekki í því að tjá hvaða útreikning eigi að framkvæma, heldur í því hvernig vélbúnaðurinn eigi nákvæmlega að gera það. Til að smíða bestu kjarnana þurfum við tungumál sem býður upp á réttar abstraksjónir yfir undirliggjandi vélbúnað. Rétta tungumálið gerir okkur kleift að tilgreina hágíðaaðgerðirnar til aðframkvæma útreikninginn, og sérhæfa þær og fínstilla af krafti á öllum punktum á breiðu vinnusvæði.

CuTeDSL er Python-bundið DSL fyrir NVIDIA GPUskjákort sem byggir ofan á CuTe skipulagsalgebru (layout algebra) og MLIR. Það þýðir rétt í tæka tíð (just-in-time) yfir í mjög fínstilltan PTX, á meðan það heldur stjórn á lágstigs vélbúnaðareiningum. Markmið þess er að bjóða upp á sömu eiginleika og afköst og CUTLASS C++, án þess dýra þýðingartíma sem heftir þróunarhraða. Af þeim tungumálum sem við höfum skoðað eða unnið með (þar á meðal CUDA, cuTile, TVM IR og Triton), teljum við það passa best við þróun kjarna okkar.

Á heildina litið, þótt það sé ekki eins hágíðalegt og kjarnyrt og Triton eða cuTile, þá komust við að því að CuTe skipulagsalgebra og samsetningar hennar (birtar í gegnum cute.Tensor) voru sanngjarnir staðgenglar. Að auki býður CuTeDSL upp á ótakmarkaðan aðgang að vélbúnaðareiningum og veitir okkur það stjórnunarstig sem við þurfum á að halda til að ná hámarksafköstum í öllum jaðartilvikum. Með CUDA voru tækifæri til sérhæfingar á þýðingartíma takmörkuð, þar sem hver viðbótarbreyta jók kostnað við þýðingartíma eða flókið byggingarkerfi um allan þjófabálk. Aftur á móti, með JIT-þýddu DSL (Just-In-Time Compiled DSL), getum við sérhæft okkur ákafari án þess að fórna endurtekningarhraða. Ólíkt TVM IR, sem notaði CUDA sem milliskref, fannst okkur MLIR-byggða pípan sem lækkaði niður í PTX vera traustari og auðveldari í villuleit en sú sem skilaði CUDA út sem milliskrefi. Að auki, meðnotkun á sérsniðnum aðgerðum til að vefja utan um PTX skipanir, komustum við að því að við gátum nýtt okkur alla eiginleika undirliggjandi vélbúnaðar til fulls.

Forritarauðvitað

Meginhluti þýðingarkostnaðar CUDA kjarna fer venjulega í sniðmátsþenslufasa (template expansion phase) kjarnanna, þar sem fínstilling og þýðing frá PTX-stigi og áfram er hægt að klára á nokkrum sekúndum. Með mörgum breytum á þýðingartíma getur þessi stigi þýðingar þanist út í margar mínútur fyrir flókinn kjarna. Þessi flöskuháls magnaður upp af því að tækifæri til samhliða þýðingar koma aðeins upp síðar, á IR- og PTX-stigum, á meðan C++ sniðmátameðhöndlun á sér stað röðunargreint (sequentially). Þótt til séu brellur til að leysa þetta með því að færa sniðmátsstofnanir yfir í aðrar þýðingareiningar, krefjast þær verulegra breytinga á byggingarkerfum og notkun á ómeðfærilegum forritaskilum (macros). Aftur á móti þurfa CuTeDSL kjarnar sem eru þýddir rétt í tæka tíð (JIT) ekki flóknar byggingarreglur og hindra ekki þýðingartíma, þar sem hægt er að beina þróunarferlinu að einum punkti á vinnusvæðinu, á meðan varkár einingarpróf síðar geta náð yfir allt hönnunarsvæði kjarnans.

Þó að JIT-þýðing færi kostnað frá þróunartíma yfir á keyrslutíma, er aukaálagið meira en ásættanlegt. Þegar líkan er tekið í notkun virkjar það aðeins fáa punkta af öllu vinnusvæðinu, sem veldur aukakostnaði upp á nokkrar sekúndur við ræsingu/upphitun ályktunarvélarinnar. Aftur á móti gæti forþýðing á öllum stillingum kjarna tekið margar mínútur. Að auki er hægt að draga enn frekar úr þessum aukakostnaði með því að skyndiminni (cache) þær JIT-þýddu afurðir.

Hinn þungi þýðingarkostnaður CUDA er magnaður upp af erfiðleikunum við að rekja uppruna villna þegar sniðmátsbundin útsending og sérhæfing er notuð, eins og raunin er með CuTe og CUTLASS. Með DSL er hægt að benda á þá kóðalínu sem myndaði villuna, sem gerir villuleit mun auðveldari. Að auki er miklu einfaldara að vensla mynduðu kóðann við DSL-skipunina sem upphaflega bjó hann til í gegnum millistig MLIR eða með því að treysta á villuleitarupplýsingar sem fylgja mynduðum PTX. Þökk sé þessum þáttum DSL fannst okkur við vera mun afkastameiri við að þróa kjarna.

Stórt vinnusvæði

Þótt kjarnar eigi margt sameiginlegt þvert á líkön, þurfa þeir samt að vera búnir ýmsum stillingarhnöppum sem geta haft áhrif á afköst. Viðeigandi breytur eru meðal annars:

Falin vídd / Höfuðvídd / Fjöldi athyglishöfuða (Hidden Dimension/Head Dimension/Number of Attention Heads)

Inntaks-/úttaksgagnagerðir (Input/Output data types)

Kvantunarkerfi (Quantization scheme)

Fastar á þýðingartíma (vægi, virkjunarföll o.s.frv.)

ROSE-arkitektúrinn, þar á meðal vélina, líkanið, kjarnasafnið og þjónustur.

Sérhæfing á þýðingartíma er gagnleg fyrir bæði afköst og villuleit. Sérstaklega í leiðsluvinnslumanni (pipelined processor), eins og raunin er með GPUskjákort, er grein sem er ekki til betri en grein sem er alltaf tekin eða aldrei tekin. Að auki skilar þýðing á kjarna fyrir ákveðið sett af stillingarhnöppum færri PTX/SASS skipunum, sem eru auðveldari í greiningu og villuleit þegar cuda-gdb er notað við rannsókn.

Fyrir tiltekna lotustærð (batch size) og stillingu má áætla væntanlega umferð í DRAM og fjöldann af reikniaðgerðum. Í kjölfarið er hægt að taka ákvarðanir frá tilviki til tilviks um ákjósanlega úthlutun á hnitneti (grid), samstillingu og notkun á ákveðnum vélbúnaðareiginleikum. CuTeDSL gerir okkur kleift að skilgreina og endurnýta íhluti fyrir djúpa sérhæfingu án þess að valda óásættanlegum aukaostnaði við þýðingu.

Sérhæfing í forskráningu og afkóðun

Almennt séð þarf að fínstilla afköst ályktunarkjarna fyrir tvö aðskilin notkunartilvik:

Forskeiddu (Prefill): lítil lotustærð (fáir hlutar, raðir), mörg tákn í röð. Hámarka þarf afköst (throughput) fyrir bæði innfellingar og inntakstákn LLM. Biðtími (latency) á bilinu tugir til hundrað millisekúndna.

Afkóðun (Decode): litlar til stórar lotustærðir, eitt (ekki forspárleg afkóðun) til nokkur tákn í röð (drög-markmiðs forspá, fjöl-tákn spá). Biðtími nálægt ræsibiðtíma kjarna.

Fyrir GEMM og athygli keyrir forskeidduáfanginn á miklum fjölda tákna og afköst ráðast af reikniafköstum tækisins. Afkóðun felur í sér sömu aðgerð á litlum lotustærðum, þar sem keyrslutímar einkennast af kostnaði við að hlaða virkjunum og vægi frá DRAM yfir í streymis-fjölvinnslumenn (streaming multiprocessors) og tensórkjarna, og eru takmarkaðir af minnisbandbreidd tækisins í staðinn. Jafnvel þótt lotustærðin eða heildarfjöldi tákna í lotu sé nokkuð skýr vísir, er skiptimörkum þröskuldurinn á milli forskeiddu og afkóðunar lúmskari vegna þess að minnisbandbreiddin er einnig ákvörðuð af faldu vídd líkansins og gagnagerð virkjunarinnar. Kjörnum er því skipt upp í forskeiddu- og afkóðunarútfærslur ef reikniritamunur eða úthlutun hnitnets/blokkar myndi nýtast öðrum hvorum þeirra.

Breytan sem er almennt frábrugðin á milli forskeiddu- og afkóðunarkjarna er fjöldi vafninga (warps) í blokk (fjöldi þráða, deilt með vafningastærðinni 32 þræðir á vafning). Æskilegt er að gera þessa breytu að fasta á þýðingartíma, þar sem hún getur hjálpað til við úthlutun skráa (registers). Í CUDA er þetta venjulega gert með því að tilgreina fjölda þráða í gegnum __launch_bounds__ og gefa upp fjölda vafninga sem sniðmátsbreytu:

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

Fjöldi skráa sem eru í boði fyrir alla þræði í blokk er vélbúnaðartakmörkun, 64K á Hopper og Blackwell. Ef fjöldi þráða er þekktur fyrir fram, veit þýðandinn nákvæmlega hversu margar skrár hann getur úthlutað hverjum þræði og getur tekist að þýða enn betur fínstillta kjarna. Hins vegar fylgir þessu kostnaður á þýðingartíma, þar sem nú þarf að þýða kjarnann aftur fyrir mismunandi vafningastærðir. Með CuTeDSL er þetta ekki lengur vandamál, þar sem við þurfum ekki að greiða kostnaðinn við þýðingu fyrir allt vinnusvæðið fyrir fram. DSL-ið gerir einnig ráð fyrir að blokkastærðin sé fastatjáning (constant expression), og framkvæmir í raun fínstillinguna sem __latch_bounds__ gefur til kynna í CUDA.

Auk þess að gera fjölda vafninga að föstu, er æskilegt að breyta honum á milli forskeiddu og afkóðunar. Til að nýta minnisbandbreidd GPUskjákorts til fulls dugar ekki alltaf að úthluta háum fjölda vafninga á blokk, þar sem aðeins er hægt að ná fullri bandbreidd með því að nota sem flesta af tiltækum streymis-fjölvinnslumönnum (SM). Því veljum við fyrst, fyrir tiltekinn fjölda tákna og þekktar gagnagerðir og faldar víddir, kjörna vixl-fyrir-þráð vigurvídd (8 fyrir bfloat16, 4 fyrir float32) til að tryggja að myndaði kóðinn noti 128-bita hleðslu-/vistunarskipanir (256-bita á Blackwell). Með því að þekkja vigurstærðina getum við reiknað út fjölda vafninga sem kjarninn þarfnast. Við reynum fyrst að hámarka fjölda SM: ef það eru færri vafningar en SM, er einn vafningur á blokk valinn. Þetta er venjulega afkóðunarafbrigði kjarnans. Ef nægir vafningar eru til staðar, aukum við fjölda vafninga smám saman í 8, 16 eða 32, og prófum efri mörk þar sem við sjáum almennt ekki lengur framför í afköstum.

Þótt Triton hafi framkvæmt sumar þessara fínstillinga undir húddinu, komust við að því að við getum náð svipuðum árangri með því að nýta CuTe skipulagsalgebru vandlega, án þess að vera takmörkuð af göllum hágíðathýðanda. Sérstaklega þegar undirvafnings láréttar lækkanir (sub-warp horizontal reductions) eiga í hlut, fannst okkur CuTeDSL bjóða upp á þá fínkornóttu stjórn sem þarf til að búa til kjörinn kóða.

Samstilling hnitnets (Grid synchronization)

Í sumum aðstæðum, til að dreifa vinnu yfir sem flestar blokkir, er hægt að nota cooperative_groups::this_grid().sync() til að samstilla allar þátttökublokkir. Takmarkanir eiga við: þar sem ekki er hægt að taka blokk úr áætlun (de-scheduled) frá SM þegar hún hefur hafið keyrslu, til að tryggja áframhaldandi framgang, má fjöldi blokka ekki fara yfir fjölda SM sem eru í boði á tækinu, til að tryggja að hvenær sem er geti allar blokkir uppfært og kannað almennu minnisstaðsetningarnar sem taka þátt í samstillingu. Þessari takmörkun er framfylgt af cudaLaunchCooperativeKernel. Kostnaðurinn við hnitnetshindrun (grid barrier) er um 3us á Hopper, auk kostnaðarins við að hreinsa út allar fyrri minnisaðgerðir í gegnum óbeina minnishindrun í sumum aðstæðum. Á heildina litið er þetta sambærilegt við aukaálagið við að skipta kjarna í tvær ræsingar, sem leiðir til hönnunarákvörðunar sem er viðkvæm fyrir afköstum. Þrátt fyrir að Programmatic Dependent Launch (PDL) geti dregið úr ræsiaukaálagi frá og með Hopper, þá er þess virði að forðast aukaálagið fyrir afkóðunarkjarna.

Ef annað vinnuálagið krefst fleiri blokka en fjölda SM, ætti að skipta kjarnanum upp og treysta á óbeina hindrun kjarnaræsinga. Þetta á að hluta til við um forskeiddu. Ef fjöldi blokka sem er kjörinn fyrir allt vinnuálag passar inn í hnitnetið, er hnitnetshindrun æskileg, eins og raunin er með afkóðunarkjarnana sem vinna á minni fjölda tákna.

Í CuTeDSL er hægt að taka þessar ákvarðanir í ræsingarkóða hýsilsmegin (host-side) og hægt er að samsetja kjarnana og flétta þeim saman við hnitnetshindranir eftir þörfum, eins og sýnt er í dæminu hér að neðan. Með svona tvíþættri stefnu er hægt að fínstilla afkóðunarkjarna til að ná lágmarksbiðtíma og hægt er að fínstilla forskeiddukjarnana fyrir hámarksafköst. Ólíkt CUDA-byggðum lausnum er hægt að skilgreina þessa rökfræði hnitmiðaðar og með minni aukaþýðingarkostnaði.

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

Í tilviki MoE all-to-all kjarna okkar, eins og lýst er hér að neðan, hjálpaði sveigjanleikinn við að smíða kjarna með eða án hnitnetshindrana og fínstilla þá á þýðingartíma okkur að ná verulegum afkastaaukningum í gegnum CuTeDSL.

CuTeDSL Kjarnar

QK Norm

Fyrsti kjarninn sem við smíðuðum í CuTeDSL var sérhæfður RMS norm kjarni, fínstilltur fyrir litlar víddir. QK norm, sem er notað í sjálf-athyglislagi (self-attention layer) nýlegra Qwen líkana, er beitt á grundvelli hvers höfuðs, á 64, 128 eða 256 stök. Ólíkt venjulegu RMS normi, sem vinnur á faldu víddum sem fara yfir þúsundir og krefst látréttrar lækkunar í gegnum bæði sameiginlegt minni (shared memory) og skrár, þarf QK norm aðeins lækkun þvert á vafninga, með þeim möguleika að þjappa mörgum höfðum saman í einn vafning. Þetta varð notkunartilvik þar sem við gátum sérhæft okkur af krafti á þýðingartíma til að finna bestu tákn-til-vafnings úthlutanir fyrir allar höfuðvíddir og gagnagerðir.

RMS norm er nokkuð einföld aðgerð, þó hún krefjist látréttrar summu til að reikna út ferning norm vigurs:

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

Þar sem athyglishöfuð eru venjulega lítil og aðgerðir fyrir hvern þráð eru mikið vigraðar (vectorized), getur stakur vafningur reiknað út ferning normið. Til að hámarka minnisbandbreidd ætti hver þráður að vinna úr 128 bitum af inntaki, sem samsvarar 8 bfloat16 gildum. Því dugar hálfur vafningur til að vinna úr höfði, og hægt er að ákvarða ferning normið með því að nota látræna lækkun sem er framkvæmd með shuffle_sync_bfly. Frá höfuðvíddinni og gagnagerðinni getum við ákvarðað fjölda höfuða á hvern vafning og tilgreint viðföngin og fjölda lækkunarumferða á tíma þýðingar.

Við getum þá smíðað sérhæfðan kjarna með tiltölulega fáum línum af kóða, með því að treysta á sjálfvirkt vigruð innlestrar-/skrifaðgerðir og nota lægra stigs arkitektúrsértækar aðgerðir fyrir warp-úppröðunina. Kjarninn ákveður hvaða höfuð (Q eða K) á að vinna á og auðkennir samsvarandi vægi sem á að beita. Hann les inn heilan vigur úr höfuðinu og leggur fyrst saman staðbundið. Í kjölfarið, miðað við fjölda þráða á hvert höfuð, er reiknað súlðurúm. Að lokum er RMS-rúmið, sjálfvirkt vigrað, ákveðið og skrifað aftur í minni. Færibreytur sem eru sértækar fyrir líkanið, svo sem WEIGHT_BIAS og EPS, eru fastafellanlegar.

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)

Borið saman við venjulegan RMS norm kjarna skilar þetta sérhæfða afbrigði 2-3x betri afköstum og biðtíma fyrir bæði forskeiddu og afkóðun.

MoE útsending/samsetning (Mixture-of-Experts Dispatch/Combine)

MoE útsending og samsetning stokka inntakstáknum yfir á sérfræðingana sem þeim er beint til, og framkvæma vægi-söfnun eftir að MoE GEMM kjarnarnir framkvæma viðeigandi fylkjamargfaldanir. Afköst eru sérstaklega mikilvæg þar sem ýmsar vörur reiða sig á MoE líkön af ýmsum stærðum. Miðað við uppsetningartegundirgreinum við mismunandi kjarna til að fínstilla:

TP=1: Staðbundinn kjarni fyrir stakt GPUskjákort. Útsending stokkar táknum til viðeigandi sérfræðinga og samsetning safnar þeim saman. Notað aðallega fyrir Qwen3 30B uppsetningar á H100.

DP=1, TP ≠ 1: Tensór-samsíða, eins-hnúts kjarni. Inntakinu er tvítekið, útsending velur táknin fyrir núverandi röð (rank), og samsetning safnar þeim saman þvert á allar raðir. Söfnun yfir annað hvort NVLink eða NVLS. Notað fyrir stóra forskeiddu.

DP ≠ 1, NVLink: Gagnasamsíða kjarni með valfrjálsri tensór-samsíðu. Hver röð sendir tákn til hverrar annarrar raðar sem hýsir sérstaka sérfræðinga. Beiningarupplýsingum er safnað saman, tákn eru flutt yfir NVIDIA NVLink. Notað fyrir gagnasamsíða H200 eða GB200 uppsetningar.

DP ≠ 1, InfiniBand: Jafningja-til-jafningja (Peer-to-Peer) gagnasamsíða og tensór-samsíða kjarni sem keyrir þvert á bæði NVLink og NVIDIA Quantum InfiniBand. Notað fyrir fjöl-hnúta Hopper uppsetningar.

Staðbundið (Local)

Staðbundni kjarninn telur fjölda tákna sem úthlutað er til hvers sérfræðings miðað við beiningatöfluna og úthlutar einstökum vísitölu til hvers tákns innan listans yfir tákn sem tilheyra hverjum sérfræðingi. Eftir beiningarskrefið eru táknin afrituð yfir í valfrjálst fyllta biðminni (padded buffers) til að senda áfram í MoE GEMM kjarna. Með því að nota staðsetningu táknanna framkvæmir samsetningarkjarninn síðan söfnun í skrám, og skrifar endanlegu táknunum í úttaksbiðminni.

Útsendingaráfanginn er breytilegur á milli forskeiddu og afkóðunar: fyrir fá tákn eru litlar beiningarupplýsingar til að lesa út, þannig að við vinnum úr því í stakri blokk og gefum merki um afganginn af hnitnetinu, þar sem einstakar blokkir afrita hluta af táknum. Táknúthlutun er ákvörðuð með því að hækkateljara atómískt í sameiginlegu minni. Þar sem forskeiddu þarf að vinna úr fleiri beiningarupplýsingum, er beiningu skipt á milli margra blokka í öðrum kjarna. Vísitölum er fyrst úthlutað staðbundið með því að nota teljara í sameiginlegu minni og er þeim síðan safnað saman þvert á allar blokkir í almennu minni til að ákvarða staðsetningu hvers tákns. Hnitnet þessa kjarna er þvingað að fjölda SM til að geta notað hnitnets-hverjar hindranir. Síðari kjarni framkvæmir afritunina, ræstur yfir miklu stærri fjölda blokka og vafninga til að nýta tiltæka minnisbandbreidd til fulls.

Samsetningarkjarninn er hlutfallslega einfaldur, þar sem hann framkvæmir vigraða söfnun, treystir á beiningarupplýsingarnar og hliðranirnar (offsets) sem útsending reiknaði út. CuTeDSL sannaði gildi sitt hér, þar sem hægt var að tjá útreikninginn í gegnum einfalda CuTe tensóra sem þýðast yfir í skilvirkan kóða fyrir allar gagnagerðir og faldar víddir.

Eins-hnúts (Singlenode)

Í tensór-samsíða uppsetningu eru beiningarupplýsingar og inntakstákn tvítekin. Hægt er að endurnýta staðbundna útsendingarkjarnann að fullu með litlum breytingum til að beina og velja aðeins þau tákn sem úthlutað er til sérfræðinganna sem staðsettir eru á núverandi röð. Með nokkrum breytum og sérhæfingu á þýðingartíma er staðbundnu útsendingarkjörnunum endurnýta.

Þar sem fjöldi raða er venjulega lægri en fjöldi sérfræðinga sem tákn er beint til, má gera ráð fyrir í meðaltilvikinu að hver tákn hafi að minnsta kosti einn af sínum tvítökum beint til hverrar raðar. Við þessar aðstæður er hægt að endurnýta staðbundna samsetningarkjarnann til að safna saman staðbundnum afritum af tákni, fylgt eftir af all-reduce til að leggja saman þvert á þátttökuraðirnar. Í tilviki forskeiddukjarnans ræsum við aðskilda kjarna, en afkóðun hins vegar sameinar báðar aðgerðirnar í eina ræsingu.

Fyrir all-reduce notum við venjulegt NVLink ef aðeins tvær raðir taka þátt, og skiptum yfir í NVLS þegar 4 eða fleiri raðir eiga hlut að máli til að fá 10-20% aukningu í afköstum. Þökk sé DSL getum við deilt megninu af lækkunarrökfræðinni og skiptum á jafningjabiðminni á milli MoE og venjulegu all-reduce kjarnanna. Þetta eru aðstæður þar sem nákvæm stjórn á vélbúnaðinum er gagnleg, þar sem DSL-aðgerðir yfir multimem.ld_reduce og multimem.st eru nauðsynlegar fyrir NVLS. Útfærsla lækkunarinnar sjálfrar fylgir að stórum hluta all_reduce dæminu úr NCCL prófunum.

Gagnasamsíða (Data Parallel)

Gagnasamsíðu kjarnarnir reiða sig á NVLink til að safna saman táknum fyrir mörg tæki. Í útsendingaráfanganum, áður en táknum er skipt út, beina jafningar fyrst staðbundið til að greina markhliðrunina á hverri röð fyrir hvert tákn sem þeir eru að senda. Í kjölfarið byrja þeir að afrita táknin yfir í móttökubiðminni sem úthlutað er fyrir hvern sendanda til að forðast árekstra. Slepp-vistunar skrif (release-store writes) fyrir hvert tákn eru notuð til að gefa merki um að lokum sé náð. Við móttökuenda er beiningarupplýsingum safnað saman til að geta þjappað öllum táknum í eitt biðminni fyrir MoE GEMM. Móttakendur kanna fána jafningjaðra skrifa með kaup-hleðslu (acquire-load) til að bíða eftir að skrifum ljúki.

Á samsetningarhliðinni eru tákn afrituð yfir NVLink á upprunaraðirnar, með fána fyrir hvert tákn sem gefur til kynna móttöku. Samsetningarkjarninn er svipaður staðbundnu og eins-hnúts afbrigðunum, með smáum mun á vistun og einhverri viðbótarsamstillingarrökfræði.

Við smíðuðum safn af PTX vöflum (wrappers) til að birta slepp-kaup hleðslur og vistanir á bæði kerfis- og tækisstigi til að útfæra þennan kjarna. Með því að bæta fínkornótt samstillingarfána og úthluta fleiri blokkum til að senda og móttaka tákn yfir NVLink, náðum við 2x framförum um fram fyrri CUDA kjarna okkar.

Jafningja-til-jafningja (Peer-to-Peer)

Jafningja-til-jafningja kjarnar okkar náðu áður fullkomnustu afköstum á Hopper GPUskjákortum sem pöruð voru við ConnectX-7 InfiniBand millistykki. Við höfum síðan flutt samsetningarkjarnana yfir í CuTeDSL fyrir frekari afkastaauka. Þar sem CUDA útgáfurnar treystu nú þegar mikið á sérhæfingu á þýðingartíma, reynðust þær vera flöskuhálsar þýðingartíma í geymslunni okkar, sem neyddi okkur til að sleppa sumum breytum. Með því að flytja yfir í JIT-þýdda DSL sérhæfðum við okkur fyrir fleiri breytur og breyttum hnitnetisúthlutunum, sem kynnti frekari mun á milli forskeiddu og afkóðunar fyrir um 3% aukningu í biðtíma.

Blackwell Stuðningur

Frá sjónarhóli ROSE kjarna fellur Blackwell stuðningur í tvo flokka: arkitektúr-óháða og arkitektúr-sértæka kjarna. Arkitektúr-sértækir kjarnar, sem reiða sig á Tensórkjarna, krefjast sérstakrar vinnu til að laga. Hins vegar, þökk sé notkun DSL, laga arkitektúr-óháðir kjarnar sig áreynslulaust.

Blackwell GPUskjákort er með allt að 160 SM, í mótsögn við 132 hjá Hopper. Fyrir flesta kjarna er þetta fasti á þýðingartíma og rökfræði í hýsilsmegin ræsingarkóða tekur mið af því við að ákvarða kjörinn kjarna (forsk. á móti afkóðun) til að keyra. Að auki kynnir Blackwell einnig 256-bita hleðsluskipanir, sem auðvelt er að senda út með því að auka vigurbreiddina (síðustu vídd CuTe tensóra) til að bæta nýtingu minnisbandbreiðar.

Blackwell kynnir vélbúnaðarstuðning fyrir verulegan fjölda kvantunarkerfa (MXFP8, MXFP4, NVFP4). Stuðningur við þetta var að mestu leyti útfærður með því að breyta núverandi kjörnum fyrir FP8 blokk-bundna skalun. JIT stuðningur er mikilvægur hér, þar sem hægt er að styðja við nýja ham án þess að þenja út stillingarsvæði þýðingartíma. Í fortíðinni fannst okkur Triton óþægilegt við að tjá látrænar lækkanir þvert á fá stök, yfir hálfan eða fjórðung af vafningi. Fínkornóttari stjórn á lækkunum var mikilvæg til að ná kjörnum afköstum þegar algildi (absolute maximum value) er reiknað yfir 16 eða 32 stök fyrir NVFP4 og MXFP8/MXFP4, í sömu röð.

Gervigreindar-aðstoðuð kjarnasmíð (AI-Assisted Kernel Generation)

Sem tiltölulega nýtt tungumál höfðu LLM litla kynningu á CuTeDSL kóða við forþjálfun sína. Að auki var stíll núverandi kóða okkar, þar sem kjarnar eru fínstilltir fyrir virkjun og norm, dálítið frábrugðinn fyrirliggjandi CUTLASS eða Quack dæmum, sem voru mikið fínstillt fyrir TMA og arkitektúr-sértækar fylkjamargföldunarskipanir. Að ná því marki að LLM gætu aðstoðað við kjarnaskrift var hægfara ferli.

Við smíðuðum handvirkt upphafssafn af kjörnum og tólföllum (utility functions), sem fylgdu stöðugum mynstrum, áður en við gátum treyst á umboðsmenn (agents) fyrir kjarnabreytingu. Þegar safnið þróaðist gátum við framkvæmt grófar þýðingar frá CUDA yfir í CuTeDSL, með verulegri handvirkri endurvinnu í kjölfarið. Sérstaklega komustum við að því að líkön áttu í erfiðleikum með breyta-til-tensors umbreytingar og vigrun, og föstu við að reyna að laga tegundarvillur (type errors) þegar stefnt var að því að fá próf til að þýðast. Þau hafa sýnt lítinn skilning eðlis hnitnets-undirstaða kjarnanna, og ruglað saman þráða- og hnitnetssamstillingu. Það sem truflaði mest var freistingin til að hugaróma (hallucinate) ótilvistar þýðingar á CUDA-sértækum bókasöfnum í stað þess að smíða gegn okkar eigin CuTeDSL skiptum. Með tímanum, eftir því sem safn CuTeDSL kjarna stækkaði, komustum við að því að LLM geta smíðað stærsta hluta beinagrindakóðans og fyllt í auðu reitina, sem gerir okkur kleift að einbeita okkur að hágíðahönnuninni.

Niðurstaða og framtíðarvinna

Borið saman við CUDA finnst okkur CuTeDSL skila miklu betri afköstum á þýðingartíma, sem gerir okkur kleift að sérhæfa okkur fyrir fleiri breytur og ná afkastaaukningu. Þótt kjarnar séu örlítið orðmargari vegna skorts á ákveðnum abstraksjónum miðað við Triton, finnst okkur flóknin vera réttlætanleg þar sem fínkornóttari stjórn á vélbúnaðinum gerir okkur kleift að smíða hraðari kjarna.

Þegar þetta er skrifað er afkastabilið á milli CUTLASS C++ og CuTeDSL kjarna lágmark, sem opnar spennandi tækifæri til að samsetja og bræða hágíða GEMM kjarna við samskipti milli tækja, stöðlun eða virkjun. Með því að skuldbinda okkur við þetta DSL munum við geta dregið frekari afköst út úr tækjunum okkar til að bæta afköst þvert á vinnuálag okkar, allt frá innfellingum til LLM.

Heimildir

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/

Tilvitnun

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