CuTeDSL hos Perplexity

Hos Perplexity bruker vi den egenutviklede kjøretidsoptimeringstjenesten Runtime-Optimized Serving Engine (ROSE) til å kjøre modeller, fra innebygginger til språkmodeller (LLM-er) med billioner av parametere. Kjøretiden vår, som ligger bak alle API-er, inkludert Sonar, Search og Embeddings, ble bygget for at forskere skal kunne utforske de best egnede modellene og arkitekturformene, samtidig som de raskt kan sette de beste kandidatene i produksjon. For å oppnå dette må vi kunne implementere, distribuere og optimere inferens for avanserte arkitekturer i høyt tempo. I dette innlegget utforsker vi hvordan CuTeDSL-kjerner integrert i inferensmotoren vår gjør det mulig å bygge de nødvendige GPU-kjernene for å bringe modeller opp til toppytelse på NVIDIA-GPU-er.

Inferens hos Perplexity AI

Tilpassede modeller, bygget av ulike team hos Perplexity, kjøres internt på NVIDIA Hopper- og Blackwell-GPU-er ved hjelp av ROSE-inferensmotoren, som er utviklet av AI Inference-teamet. ROSE ble opprinnelig bygget for å betjene tilpassede Llama-modeller bedre for både språkodekoding og klassifisering, med et grensesnitt som er kompatibelt med NVIDIA Triton Inference Server. Over tid utviklet den seg til en fullverdig inferensmotor som betjener betydelige LLM-er, samt rangerings-, klassifiserings-, poengsettings- og andre transformerbaserte modeller for et bredde av team som driver Perplexitys produkter. Designtålet for ROSE er å være enkelt: Over tid investeres det innsats for å oppnå toppytelse, men førsteprioriteten er å tilpasse seg de stadig evig utviklende arkitekturene og bruksområdene for modellene våre.

Forstørr

På et overordnet nivå består ROSE av en motor som tilpasser en modell til et grensesnitt som klienter kan samhandle med. Motorene kan gjenbrukes i stor grad, og håndterer forespørselsplanlegging og -bunting, samt henting av avkodede tokener, innebygginger eller poengsummer fra de underliggende modellene. Initialiseringen av enheter, kommunikasjonskanaler mellom prosesser og noder, og vektlasting håndteres av dem. For LLM-er er motoren ansvarlig for bunting, oppdeling i biter (chunking) og sampling. Den er også ansvarlig for tildeling av KV-lagring, for både full og lineær oppmerksomhet, samt prefikssamsvar. Innebygde motorer er enklere, da de bare utfører nettbasert bunting. Utførelsen er optimert slik at motoren forbereder neste runde med arbeid mens modellen kjører parallelt på akseleratoren.

ROSE eksponerer en rik samling av egendefinerte lag for modellering, og omslutter GPU-kjerner bygget med CuTeDSL, Triton, CUDA, CUTLASS og cuBLAS. For matrisemultiplikasjon og oppmerksomhetsimplementeringer baserer vi oss primært på NVIDIA-kjerner fra CUTLASS og cuBLAS. Resten av rørledningen (pipeline), inkludert innebygging, norm, MoE og aktiveringskjerner, er en del av kjernebiblioteket til ROSE. For best mulig å støtte maskinvaren vi kjører inferens på, er ROSE utstyrt med egne enhetskommunikasjonsprimitiver for MoE-rutting, utsending og kombinering. Selv om de andre kjernene er relativt enkle, må de spesialiseres grundig for å oppnå toppytelse på alle de forskjellige skjulte dimensjonene og enhetene de kjører på. I tillegg baserer vi oss på kompileringstidspesialisering for å imøtekomme alle innstillinger og justeringer for ulike modeller (vektbias vs. ingen vektbias for RMS-norm, aktiveringsfunksjoner osv.). Derfor bestemte vi oss for å bruke DSL-er til å implementere dem: vi bygget opprinnelig ROSE med Triton. Siden den gang har vi imidlertid tatt i bruk CuTeDSL som det primære miljøet for GPU-programmering.

Hvorfor CuTeDSL?

Fra et matematisk perspektiv er inferenskjerner vanligvis ukompliserte. Følgelig ligger kompleksiteten ved utvikling av en kjerne generelt ikke i å uttrykke hvilken beregning som skal utføres, men i nøyaktig hvordan maskinvaren skal gjøre det. For å bygge de beste kjernene trenger vi et språk som tilbyr de riktige abstraksjonene over den underliggende maskinvaren. Det riktige språket gjør oss i stand til å spesifisere operasjonene på høyt nivå for å utføre beregningen, og aggressivt spesialisere og optimere dem på tvers av alle punkter i et bredt konfigurasjonsrom.

CuTeDSL er et Python-basert DSL for NVIDIA-GPU-er bygget oppå CuTe-oppsettalgebra og MLIR. Det kompileres just-in-time til svært optimert PTX, samtidig som det beholder kontrollen over lavnivå maskinvareprimitiver. Målet er å tilby de samme funksjonene og ytelsen som CUTLASS C++, uten de dyre kompileringstidene som hemmer utviklingshastigheten. Blant språkene vi har undersøkt eller jobbet med (inkludert CUDA, cuTile, TVM IR og Triton), anser vi det for å være det beste alternativet for utvikling av kjernene våre.

Totalt sett, selv om det ikke er like høynivå og konsist som Triton eller cuTile, fant vi ut at CuTe Layout Algebra og dens sammensetninger (eksponert via cute.Tensor) var rimelige erstatninger. I tillegg tilbyr CuTeDSL ubegrenset tilgang til maskinvareprimitiver, noe som gir oss den kontrollen vi trenger for å oppnå toppytelse i alle spesialtilfeller (edge cases). Med CUDA var mulighetene for kompileringstidspesialisering begrenset, ettersom hver ekstra parameter la til en urimelig kostnad for kompileringstidene eller kompleksitet i byggesystemet. I kontrast kan vi med en Just-In-Time-kompilert DSL spesialisere mer aggressivt uten å offerere iterasjonstempoet. I motsetning til TVM IR, som brukte CUDA som et mellomtrinn, fant vi at den MLIR-baserte rørledningen som senket til PTX var mer robust og enklere å feilsøke enn en som sendte ut CUDA som et mellomtrinn. I tillegg fant vi ved bruk av egendefinerte operasjoner (custom ops) til å omslutte PTX-instruksjoner, at vi fullt ut kunne utnytte alle funksjonene til den underliggende maskinvaren.

Utvikleropplevelse

Hoveddelen av kompileringskostnaden for CUDA-kjerner brukes vanligvis i malekspansjonsfasen (template expansion) for kjernene, ettersom optimering og kompilering fra PTX-stadiet og utover kan fullføres i løpet av få sekunder. Med mange kompileringstidsparametere kan dette kompileringsstadiet eksplodere til flere minutter for en ikke-triviell kjerne. Denne flaskehalsen forsterkes av det faktum at muligheter for parallell kompilering først oppstår senere, på IR- og PTX-nivå, der C++-malhåndtering skjer sekvensielt. Selv om det finnes triks for å løse dette ved å flytte malinstansieringer til separate kompileringsenheter, krever de betydelige justeringer av byggesystemer og bruk av uhåndterlige makroer. I kontrast krever ikke CuTeDSL-kjerner som kompileres Just-In-Time (JIT) komplekse bygningsregler og hemmer ikke kompileringstidene, ettersom utviklingssyklusen kan fokusere på et enkelt punkt i konfigurasjonsrommet, mens grundige enhetstester senere kan dekke hele designrommet for kjernen.

Selv om JIT-kompilering flytter kostnader fra utviklingstid til kjøretid, er overheaden mer enn akseptabel. Når en modell distribueres, aktiverer den bare noen få av punktene fra hele konfigurasjonsrommet, noe som medfører overhead på sekunders skala under oppstart/oppvarming av inferensmotoren. I kontrast kan forhåndskompilering av alle konfigurasjoner for en kjerne strekke seg over flere minutter. I tillegg kan overhead dempes ytterligere ved å hurtigbufre de JIT-kompilerte artefaktene.

Den tunge kompileringskostnaden for CUDA forsterkes av vanskeligheten med å spore opprinnelsen til feil når malbasert utsending og spesialisering brukes, slik tilfellet er med CuTe og CUTLASS. Med en DSL kan feil pekes direkte på den kodelinjen som genererte dem, noe som gjør feilsøking betydelig enklere. I tillegg er det mye enklere å tilordne den genererte koden til DSL-setningen som genererte den, via mellomliggende MLIR-stadier eller ved å stole på feilsøkingsinformasjon knyttet til den genererte PTX-en. Takket være disse aspektene ved DSL-en fant vi oss selv mye mer produktive under utvikling av kjerner.

Stort konfigurasjonsrom

Selv om kjerner deler mange likheter på tvers av modeller, må de likevel utstyres med en rekke konfigurasjonsknapper som kan ha innvirkning på ytelsen. Relevante parametere inkluderer:

Skjult dimensjon / Hodedimensjon / Antall oppmerksomhetshoder

Inndata-/utdatadatatyper

Kvantiseringsskjema

Kompileringstidskonstanter (bias, aktiveringsfunksjoner osv.)

Forstørr

Kompileringstidspesialisering er nyttig for både ytelse og feilsøking. Særlig i en rørledningsprosessor (pipelined processor), som tilfellet er med GPU-er, er en ikke-eksisterende gren bedre enn en gren som alltid blir tatt eller aldri blir tatt. I tillegg resulterer kompilering av en kjerne for et bestemt sett med konfigurasjonsknapper i færre PTX-/SASS-instruksjoner, som er enklere å analysere og feilsøke ved bruk av cuda-gdb under en undersøkelse.

For en gitt batch-størrelse og konfigurasjon kan den forventede trafikken til DRAM og antall aritmetiske operasjoner estimeres. Deretter kan det tas optimale beslutninger fra sak til sak for optimal rutetildeling (grid allocation), synkronisering og bruk av visse maskinvarefunksjoner. CuTeDSL lar oss definere og gjenbruke komponenter for dyp spesialisering uten å pådra oss en uakseptabel overhead i kompileringstid.

Forhåndsutfylling og dekodingsspesialisering

Generelt må ytelsen til inferenskjerner optimeres for to separate bruksområder:

Prefill: liten batch-størrelse (få biter, sekvenser), mange tokener per sekvens. Gjennomstrømning skal maksimeres for både innebygginger og inndatatokenene til LLM-er. Latens på ti til hundrevis av millisekunder.

Avkoding (Decode): liten til stor batch-størrelse, ett (ingen spekulativ avkoding) til noen få tokener per sekvens (utkast-mål-spekulasjon, flertokent预测). Latens nær kjernelanseringslatens.

For GEMM og oppmerksomhet (attention) kjører prefill-stadiet på et stort antall tokener, og ytelsen er begrenset av enhetens beregningsgjennomstrømning. Avkoding innebærer den samme operasjonen på små batch-størrelser, der utførelsestidene domineres av kostnaden ved å laste aktiveringene og vektene fra DRAM inn i strømningsmultiprosessorene og tensor-kjernene, begrenset av enhetens minnebåndbredde i stedet. Selv om batch-størrelsen eller det totale antallet tokener i en batch er en relativt tydelig indikator, er overgangsterskelen mellom prefill og avkoding mer subtil fordi minnebåndbredden også bestemmes av modellens skjulte dimensjon og aktiveringens datatype. Kjernene deles derfor opp i prefill- og avkodingsimplementeringer dersom algoritmiske forskjeller eller rute-/blokktildelinger vil være til fordel for en av dem.

Parameteren som universelt skiller seg mellom prefill- og avkodingskjerner, er antall warp-er i en blokk (antall tråder delt på warp-størrelsen på 32 tråder per warp). Det er ønskelig å gjøre denne parameteren til en kompileringstidskonstant, da det kan bidra til registertildeling. I CUDA gjøres dette typisk ved å spesifisere antall tråder via __launch_bounds__ og oppgi antall warp-er som en malparameter:

Antall registre tilgjengelig for alle tråder i en blokk er en maskinvarebegrensning, 64K på Hopper og Blackwell. Hvis antall tråder er kjent på forhånd, vet kompilatoren nøyaktig hvor mange registre den kan tilordne til hver tråd, og kan lykkes med å kompilere mer aggressivt optimerte kjerner. Dette kommer imidlertid med en kostnad i kompileringstid, ettersom kjernen nå må rekompileres for forskjellige warp-størrelser. Med CuTeDSL er dette ikke lenger et problem, siden vi ikke trenger å betale kompileringskostnaden for hele konfigurasjonsrommet på forhånd. DSL-en forventer også at blokkstørrelsen er et konstant uttrykk, og utfører implisitt den optimeringen som antydes av __latch_bounds__ i CUDA.

I tillegg til å gjøre antall warp-er til en konstant, er det ønskelig å variere det mellom prefill og avkoding. For å utnytte minnebåndbredden til en GPU fullt ut er det ikke tilstrekkelig alltid å tilordne et høyt antall warp-er per blokk, siden full båndbredde bare kan oppnås ved å bruke så mange som mulig av de tilgjengelige strømningsmultiprosessorene (SM-ene). Følgelig velger vi først en ideell vektordimensjon per tråd (8 for bfloat16, 4 for float32) for et gitt antall tokener og kjente datatyper og skjulte dimensjoner, for å sikre at den genererte koden bruker 128-biters laste-/lagringsinstruksjoner (256-biters på Blackwell). Når vi kjenner vektorstørrelsen, kan vi beregne antall warp-er som kreves av kjernen. Vi prøver først å maksimere antall SM-er: hvis det er færre warp-er enn SM-er, velges én enkelt warp per blokk. Dette er typisk avkodingsvarianten av kjernen. Hvis det er nok warp-er, øker vi gradvis antall warp-er til 8, 16 eller 32, og kjører benchmarking for en øvre grense som gjør at vi vanligvis ikke lenger ser noen forbedring i ytelse.

Selv om Triton utførte noen av disse optimeringene under panseret, fant vi ut at vi kan oppnå lignende resultater ved forsiktig bruk av CuTe-oppsettalgebra, uten å bli begrenset av manglene til en høynivåkompilator. Særlig når det er involvert under-warp-horisontale reduksjoner, fant vi at CuTeDSL ga oss den finkornede kontrollen som kreves for å generere optimal kode.

Rutesynkronisering (Grid synchronization)

I noen situasjoner, for å distribuere arbeid på tvers av flest mulig blokker, kan cooperative_groups::this_grid().sync() brukes til å synkronisere alle deltakende blokker. Restriksjoner gjelder: Siden en blokk ikke kan avplanlegges fra en SM når den først har startet utførelsen, kan ikke antall blokker overskride antall SM-er som er tilgjengelige på enheten, for å garantere fremdrift og sikre at alle blokker til enhver tid kan oppdatere og avspørre (poll) de globale minneplassene som er involvert i synkronisering. Denne restriksjonen håndheves av cudaLaunchCooperativeKernel. Kostnaden for en rutesperre (grid barrier) er ca. 3 us på Hopper, pluss kostnaden ved å tømme alle tidligere minneoperasjoner via en implisitt minnesperre i noen situasjoner. Totalt sett kan dette sammenlignes med overheaden ved å dele en kjerne inn i to lanseringer, noe som fører en ytelsessensitiv designtakst. Til tross for at Programmatic Dependent Launch (PDL) kan redusere lanseringskostnader fra og med Hopper, er det verdt å unngå denne overheaden for avkodingskjerner.

Hvis en av arbeidsmengdene krever flere blokker enn antall SM-er, bør kjernen deles opp ved å stole på den implisitte sperren for kjernelanseringer. Dette gjelder primært for prefill. Hvis antall blokker som er ideelt for alle arbeidsmengder, passer inn i ruten, foretrekkes en rutesperre, slik tilfellet er med avkodingskjernene som opererer på et mindre antall tokener.

I CuTeDSL kan disse beslutningene tas i lanseringskoden på vertsiden, og kjernene kan sammensettes og flettes med rutesperrer etter behov, som illustrert i eksemplet nedenfor. Med en slik hybrid strategi kan avkodingskjerner optimeres for å nå minimal latens, og prefill-kjerner kan optimeres for maksimal gjennomstrømning. I motsetning til CUDA-basede løsninger kan denne logikken defineres mer konsist og med mindre overhead i kompileringstid.

Når det gjelder våre MoE all-to-all-kjerner, som beskrevet nedenfor, hjalp fleksibiliteten til å bygge kjerner med eller uten rutesperrer og optimere dem ved kompileringstid oss med å oppnå betydelige ytelsesgevinster gjennom CuTeDSL.

CuTeDSL-kjerner

QK-norm

Den første kjernen vi bygget i CuTeDSL, var en spesialisert RMS-normkjerne, optimert for små dimensjoner. QK-norm, som brukes i selvoppmerksomhetslaget (self-attention) i de nyere Qwen-modellene, anvendes på hode-for-hode-basis (per-head), på 64, 128 eller 256 elementer. I motsetning til vanlig RMS-norm, som opererer på skjulte dimensjoner som overstiger tusenvis og krever en horisontal reduksjon gjennom både delt minne og registre, trenger QK-norm bare reduksjon på tvers av warp-er, med mulighet for å presse flere hoder inn i en enkelt warp. Dette ble et bruksområde der vi kunne optimere aggressivt ved kompileringstid for å finne de beste token-til-warp-tilordningene for alle hodedimensjoner og datatyper.

RMS-norm er en ganske rett fram operasjon, selv om den krever en horisontal sum for å beregne den kvadrerte normen til en vektor:

Siden oppmerksomhetshoder vanligvis er små og trådroperasjoner er sterkt vektoriserte, kan en enkelt warp beregne kvadratnormen. For å maksimere minnebåndbredden bør hver tråd behandle 128 bit med inndata, noe som tilsvarer 8 bfloat16-verdier. Følgelig er en halv-warp (half-warp) tilstrekkelig til å behandle et hode, og den kvadrerte normen kan bestemmes ved hjelp av en horisontal reduksjon utført av shuffle_sync_bfly. Med utgangspunkt i hodedimensjonen og datatypen kan vi bestemme antall hoder per warp og spesifisere argumentene og antall reduksjonsrunder på kompileringstidspunktet.

Vi kan deretter bygge en spesialisert kjerne med relativt få kodelinjer, ved å stole på automatisk vektoriserte innlesinger/lagringer og bruke lavere-nivå arkitekturspesifikke operasjoner for warp-reduksjonen. Kjernen bestemmer hvilket hode (Q eller K) den skal arbeide med, og identifiserer de tilsvarende vektene som skal brukes. Den laster inn en hel vektor fra hodet og summerer først opp lokalt. Deretter beregnes den kvadrerte normen basert på antall tråder per hode. Til slutt bestemmes RMS-normen, automatisk vektorisert, og skrives tilbake til minnet. Modellspesifikke parametere, som WEIGHT_BIAS og EPS, blir konstantfoldet (constant-folded).

Sammenlignet med en vanlig RMS-normkjerne leverer denne spesialiserte varianten 2–3x bedre gjennomstrømning og latens for både prefill og avkoding.

Mixture-of-Experts-utsending/kombinering (Dispatch/Combine)

MoE-utsending og -kombinering stokker om inndatatokener til de ekspertane de rutes til, og utfører en vektet akkumulering etter at MoE GEMM-kjerner utfører de relevante matrisemultiplikasjonene. Ytelse er av spesiell interesse, ettersom ulike produkter er avhengige av MoE-modeller i ulike størrelser. Basert på distribusjonstypene skiller vi mellom forskjellige kjerner som skal optimeres:

TP=1: Lokal kjerne for en enkelt GPU. Utsending stokker om tokener til de relevante ekspertane, og kombinering akkumulerer dem. Brukes primært for Qwen3 30B-distribusjoner på H100.

DP=1, TP ≠ 1: Tensor-parallelle, enkeltnoders kjerner. Inndataene replikeres, utsending velger ut tokenene for gjeldende rangering (rank), og kombinering akkumulerer dem på tvers av alle rangeringer. Akkumulering på tvers av enten NVLink eller NVLS. Brukes for store prefill-enheter.

DP ≠ 1, NVLink: Dataparallell kjerne med valgfri tensor-parallelleffekter. Hver rangering sender tokener til hverandre som er vert for spesifikke eksperter. Rutinginformasjon aggregeres, og tokener overføres over NVIDIA NVLink. Brukes for dataparallelle H200- eller GB200-distribusjoner.

DP ≠ 1, InfiniBand: Peer-to-Peer dataparallell og tensor-parallell kjerne som opererer på tvers av både NVLink og NVIDIA Quantum InfiniBand. Brukes for flernoders Hopper-distribusjoner.

Lokal

Den lokale kjernen teller antall tokener som er tilordnet til hver ekspert basert på rutestabellen, og tilordner en unik indeks til hvert token i listen over tokener som tilhører hver ekspert. Etter rutingtrinnet kopieres tokenene inn i buffere som kan fylles ut (padded) for å bli sendt videre til en MoE GEMM-kjerne. Ved hjelp av tokenplasseringene utfører kombineringskjernen deretter en akkumulering i registre, og skriver de endelige tokenene til en utdatabuffer.

Utsendingsfasen varierer mellom prefill og avkoding: For få tokener er det lite rutinginformasjon å lese ut, så vi behandler den i en enkelt blokk og signaliserer resten av ruten, der individuelle blokker kopierer deler av tokener. Tokentildeling bestemmes ved å inkrementere tellere atomisk i delt minne. Siden prefill må behandle mer rutinginformasjon, deles rutingen opp på tvers av flere blokker i en separat kjerne. Indekser tildeles først lokalt ved hjelp av tellere i delt minne, og aggregeres deretter på tvers av alle blokker i global minne for å bestemme posisjonen til hvert token. Ruten for denne kjernen klemmes (clamped) til antall SM-er for å kunne bruke ruteomfattende sperrer. En etterfølgende kjerne utfører kopieringene, lansert over et mye større antall blokker og warp-er for å utnytte den tilgjengelige minnebåndbredden fullt ut.

Kombineringskjernen er relativt enkel, da den utfører vektorisert akkumulering basert på rutinginformasjonen og forskyvningene (offsets) som er beregnet av utsendingen. CuTeDSL viste seg å være verdifull her, ettersom beregningen kunne uttrykkes gjennom enkle CuTe-tensorer som kompileres til effektiv kode for alle datatyper og skjulte dimensjoner.

Enkeltnode

I en tensor-parallell distribusjon replikeres rutinginformasjonen og inndatatokenene. Den lokale utsendingskjernen kan gjenbrukes fullt ut med en liten endring for å rute og bare velge ut de tokenene som er tilordnet ekspertane som befinner seg på gjeldende rangering. Med noen parametere og kompileringstidspesialisering gjenbrukes de lokale utsendingskjernene.

Siden antall rangeringer vanligvis er lavere enn antall eksperter et token rutes til, kan man i gjennomsnitt anta at hvert token har minst én av sine replikaer rutet til hver rangering. I denne situasjonen kan den lokale kombineringskjernen gjenbrukes til å aggregere de lokale kopiene av et token, etterfulgt av en all-reduce for å summere opp på tvers av de deltakende rangeringene. Når det gjelder prefill-kjernen, lanserer vi separate kjerner, men avkodingen fletter begge operasjonene sammen til én enkelt lansering.

For all-reduce bruker vi ren NVLink hvis bare to rangeringer deltar, og bytter til NVLS når 4 eller flere rangeringer er involvert, for en økning i gjennomstrømning på 10–20 %. Takket være DSL-en kan vi dele det meste av reduksjonslogikken og peer-bufferutvekslingene mellom MoE- og de vanlige all-reduce-kjernene. Dette er en situasjon der tett kontroll over maskinvaren er gunstig, ettersom DSL-operasjoner over multimem.ld_reduce og multimem.st kreves for NVLS. Implementeringen av selve reduksjonen følger i stor grad all_reduce-eksemplet fra NCCL-tester.

Dataparallell

De dataparallelle kjernene baserer seg på NVLink for å aggregere tokener for flere enheter. I utsendingsfasen, før tokenutveksling, ruter peers først lokalt for å identifisere målforskyvningen på hver rangering for hvert token de sender. Deretter begynner de å kopiere tokenene over til mottaksbuffere som er tildelt for hver avsender for å unngå konflikter. Skrivinger med frigjørings- og lagringssemantikk per token (release-store writes) brukes til å signalisere fullføring. På mottakersiden aggregeres rutinginformasjon for å kunne kondensere alle tokener til en enkelt buffer for MoE GEMM. Mottakere avspør (poll) det peer-skrevne flagget med en innhentings- og innlastingsoperasjon (acquire-load) for å vente på at skrivingene skal fullføres.

På kombineringssiden kopieres tokener over NVLink til kilderangeringene, med et flagg per token som indikerer mottak. Kombineringskjernen ligner på de lokale og enkeltnodenes varianter, med små forskjeller i adressering og litt ekstra synkroniseringslogikk.

Vi bygget et bibliotek med PTX-omsluttere for å eksponere innhentings- og frigjøringsinnlastinger og -lagringer (release-acquire loads and stores) med omfang på både system- og enhetsnivå for å implementere denne kjernen. Ved å forbedre granulariteten til synkroniseringsflaggene og tildele flere blokker til å sende og motta tokener over NVLink, oppnådde vi 2x forbedringer over våre tidligere CUDA-kjerner.

Peer-to-Peer

Våre Peer-to-Peer-kjerner oppnådde tidligere toppytelse på Hopper-GPU-er kombinert med ConnectX-7 InfiniBand-adaptere. Vi har siden portet kombineringskjernene over til CuTeDSL for ytterligere ytelsesgevinster. Siden CUDA-versjonene allerede stolte tungt på kompileringstidspesialisering, viste de seg å være flaskehalsene for kompileringstid i repositoriet vårt, noe som tvang oss til å utelate noen parametere. Ved å migrere til den JIT-kompilerte DSL-en spesialiserte vi for flere parametere og justerte rutetildelinger, noe som introduserte ytterligere forskjeller mellom prefill og avkoding for en økning i latens på omtrent 3 %.

Blackwell-støtte

Fra perspektivet til ROSE-kjerner faller Blackwell-støtte inn i to kategorier: arkitekturuavhengige og arkitekturspesifikke kjerner. Arkitekturspesifikke kjerner, som stoler på tensor-kjernere, krever eksplisitt arbeid for å tilpasses. Takket være bruken av DSL-en tilpasser imidlertid de arkitekturuavhengige kjernene seg uten videre.

En Blackwell-GPU har opptil 160 SM-er, i kontrast til Hoppers 132. For de fleste kjerner er dette en kompileringstidskonstant, og logikken i lanseringskoden på vertsiden tar hensyn til dette for å bestemme den optimale kjernen (prefill vs. avkoding) som skal kjøres. I tillegg introduserer Blackwell også 256-biters lasteinstruksjoner, som enkelt kan sendes ut ved å øke vektorvidden (siste dimensjon i CuTe-tensorer) for å forbedre utnyttelsen av minnebåndbredden.

Blackwell introduserer maskinvarestøtte for et betydelig antall kvantiseringsskjemaer (MXFP8, MXFP4, NVFP4). Støtte for disse ble implementert primært ved å finjustere eksisterende kjerner for FP8 blokkvis skalering. JIT-støtte er avgjørende her, siden nye moduser kan støttes uten at kompileringens konfigurasjonsrom eksploderer. Tidligere syntes vi Triton var uheldig når det gjelder å uttrykke horisontale reduksjoner på tvers av få elementer, over en halv eller en kvart warp. Finkornet kontroll over reduksjoner var avgjørende for å oppnå optimal ytelse ved beregning av en absolutt maksimumsverdi på tvers av henholdsvis 16 eller 32 elementer for NVFP4 og MXFP8/MXFP4.

AI-assistert kjernergenerering

Som et relativt nytt språk hadde språkmodeller (LLM-er) liten eksponering for CuTeDSL-kode under forhåndstreningen sin. I tillegg skilte stilen på vår nåværende kode, med kjerner optimert for aktiveringer og normer, seg litt fra de eksisterende CUTLASS- eller Quack-eksemplene, som var sterkt optimert for TMA og arkitekturspesifikke matrisemultiplikasjonsinstruksjoner. Å nå et punkt der LLM-er kunne hjelpe til med kjerne-skriving, var en gradvis prosess.

Vi bygget manuelt opp et initialt korpus av kjerner og verktøyfunksjoner som fulgte konsistente mønstre, før vi kunne stole på agenter for kjernekonvertering. Etter hvert som korpuset utviklet seg, kunne vi utføre grove oversettelser fra CUDA til CuTeDSL, med betydelig manuelt etterarbeid som fulgte. Spesielt fant vi ut at modellene slet med pekertil-tensor-konverteringer og vektorisering, og ble sittende fast i forsøk på å fikse feil med datatyper når målet var å få en test til å kompilere. De har vist liten forståelse for den rutebaserte naturen til kjernene, og forveksler tråd- og rutesynkronisering. Mest frustrerende var det en fristelse til å hallusinere ikke-eksisterende oversettelser av CUDA-spesifikke biblioteker i stedet for å bygge opp mot våre egne CuTeDSL-erstatninger. Over tid, etter hvert som korpuset av CuTeDSL-kjerner vokste, fant vi ut at LLM-er kan bygge det meste av skjelettkoden og fylle ut de blanke feltene, slik at vi kan fokusere på overordnet design.

Konklusjon og fremtidig arbeid

Sammenlignet med CUDA finner vi at CuTeDSL leverer mye bedre kompileringstidsytelse, noe som gjør at vi kan spesialisere for flere parametere og oppnå ytelsesforbedringer. Selv om kjernene er litt mer ordrike (verbose) på grunn av mangelen på visse abstraksjoner sammenlignet med Triton, finner vi at kompleksiteten er berettiget siden finkonstruert kontroll over maskinvaren lar oss bygge raskere kjerner.

I skrivende stund er ytelsesgapet mellom CUTLASS C++- og CuTeDSL-kjerner minimalt, noe som åpner opp for spennende muligheter til å sammensette og flette høyytelses GEMM-kjerner med enhetskommunikasjon, normalisering eller aktiveringer. Ved å forplikte oss til dette DSL-et vil vi kunne hente ut ytterligere ytelse fra enhetene våre for å forbedre gjennomstrømningen på tvers av arbeidsmengdene våre, fra innebygginger til LLM-er.

Referanser

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/

Sitat

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