CuTeDSL hos Perplexity

Hos Perplexity bruger vi vores egen interne Runtime-Optimized Serving Engine (ROSE) til at køre modeller lige fra embeddings til LLM'er med billioner af parametre. Vores kørselsmiljø, som ligger bag alle API'er, herunder Sonar, Search og Embeddings, blev bygget til at give forskere mulighed for at udforske de bedst egnede modeller og arkitekturer og samtidig hurtigt implementere de bedste kandidater i produktion. For at opnå dette skal vi være i stand til at implementere, udrulle og optimere inferens til avancerede arkitekturer i et højt tempo. I dette blogindlæg undersøger vi, hvordan CuTeDSL-kerner, der er integreret i vores inferensmotor, gør os i stand til at bygge de nødvendige GPU-kerner for at få modeller op på topydelse på NVIDIA-GPU'er.

Inferens hos Perplexity AI

Tilpassede modeller, som er bygget af forskellige teams hos Perplexity, hostes internt på NVIDIA Hopper- og Blackwell-GPU'er ved hjælp af ROSE-inferensmotoren, der er udviklet af AI Inference-teamet. ROSE blev oprindeligt bygget til bedre at understøtte tilpassede Llama-modeller til både sprogdekodning og klassificering med en grænseflade, der er kompatibel med NVIDIA Triton Inference Server. Over tid udviklede den sig til en fuldbyrdet inferensmotor, der betjener betydelige LLM'er samt rangerings-, klassificerings-, scorings- og andre transformerbaserede modeller for en bred vifte af teams til at drive Perplexitys produkter. Designmålet for ROSE er at være enkelt: Over tid investeres der i at nå topydelse, men hovedprioriteten er at tilpasse sig de stadigt udviklende arkitekturer og use cases for vores modeller.

Udvid

Overordnet set består ROSE af en motor, som tilpasser en model til en grænseflade, som klienter kan interagere med. Motorerne er i høj grad genbrugelige, håndterer anmodningsplanlægning og batching samt henter dekodede tokens, embeddings eller scores fra de underliggende modeller. Initialiseringen af enheder, inter-proces- og inter-node-kommunikationskanaler og vægtindlæsning håndteres alt sammen af dem. For LLM'er er motoren ansvarlig for batching, chunking og sampling. Den er også ansvarlig for allokering af KV-lagring for både fuld og lineær attention samt præfiksmatching. Embedding-motorer er enklere, da de kun udfører on-line batching. Eksekveringen er optimeret således, at motoren forbereder den næste runde af arbejde, mens modellen kører parallelt på acceleratoren.

ROSE stiller en rig samling af tilpassede lag til rådighed til modellering, som indkapsler GPU-kerner bygget med CuTeDSL, Triton, CUDA, CUTLASS og cuBLAS. Til matrixmultiplikation og attention-implementeringer stoler vi primært på NVIDIA-kerner fra CUTLASS og cuBLAS. Resten af pipelinen, herunder embedding-, norm-, MoE- og aktiveringskerner, er en del af ROSE's kernebibliotek. For bedst muligt at understøtte den hardware, vi kører inferens på, er ROSE udstyret med sine egne inter-device-kommunikationsprimitiver til MoE-routing, -dispatch og -combine. Selvom de andre kerner er relativt enkle, skal de være grundigt specialiserede for at opnå topydelse på alle de forskellige skjulte dimensioner og enheder, de kører på. Derudover stoler vi på compile-time-specialisering for at imødekomme alle knapper og justeringer i forskellige modeller (vægtbias vs. ingen vægtbias for RMS-norm, aktiveringsfunktioner osv.). Derfor besluttede vi at bruge DSL'er til at implementere dem: Vi byggede oprindeligt ROSE ved hjælp af Triton. Vi har dog siden taget CuTeDSL til os som det primære miljø til GPU-programmering.

Hvorfor CuTeDSL?

Fra et matematisk perspektiv er inferenskerner sædvanligvis ukomplicerede. Når man udvikler en kerne, ligger kompleksiteten derfor generelt ikke i at udtrykke, hvilken beregning der skal udføres, men i nøjagtigt hvordan hardwaren skal gøre det. For at bygge de bedste kerner har vi brug for et sprog, der tilbyder de rette abstraktioner over den underliggende hardware. Det rette sprog sætter os i stand til at specificere de overordnede operationer til at udføre beregningen og aggressivt specialisere og optimere dem på tværs af alle punkter i et bredt konfigurationsrum.

CuTeDSL er et Python-baseret DSL til NVIDIA-GPU'er bygget oven på CuTe layout algebra og MLIR. Det kompileres just-in-time til stærkt optimeret PTX, samtidig med at kontrollen over low-level hardwareprimitiver bevares. Det sigter mod at tilbyde de samme funktioner og den samme ydeevne som CUTLASS C++ uden de dyre kompileringstider, der hæmmer udviklingshastigheden. Ud af de sprog, vi har undersøgt eller arbejdet med (herunder CUDA, cuTile, TVM IR og Triton), anser vi det for at være det bedst egnede til udviklingen af vores kerner.

Overordnet set, selvom det ikke er lige så højt niveau og kortfattet som Triton eller cuTile, fandt vi ud af, at CuTe Layout Algebra og dens sammensætninger (eksponeret via cute.Tensor) var rimelige erstatninger. Derudover tilbyder CuTeDSL ubegrænset adgang til hardwareprimitiver, hvilket giver os den grad af kontrol, vi har brug for for at opnå topydelse i alle grænsetilfælde. Med CUDA var mulighederne for compile-time-specialisering begrænsede, da hver ekstra parameter tilføjede en urimelig omkostning til kompileringstiderne eller kompleksitet til bygge-systemet. I modsætning hertil kan vi med et Just-In-Time-kompileret DSL specialisere os mere aggressivt uden at offere iterationshastigheden. I modsætning til TVM IR, som brugte CUDA som et mellemtrin, fandt vi den MLIR-baserede pipeline, der sænkede til PTX, mere robust og lettere at feilsøge end en, der udsendte CUDA som et mellemtrin. Derudover fandt vi ud af, at vi med brugen af tilpassede ops til at indkapsle PTX-instruktioner fuldt ud kunne udnytte alle funktionerne i den underliggende hardware.

Udvikleroplevelse

Hovedparten af kompileringsomkostningerne for CUDA-kerner bruges typisk i skabelonudvidelsesfasen for kerner, da optimering og kompilering fra PTX-stadiet og fremefter kan udføres på få sekunder. Med mange compile-time-parametre kan denne fase af kompilering eksplodere til flere minutter for en ikke-triviel kerne. Denne flaskehals forstærkes af det faktum, at muligheder for parallel kompilering først opstår senere på IR- og PTX-niveau, hvor C++-skabelonhåndtering sker sekventielt. Selvom der findes tricks til at imødegå dette ved at flytte skabeloninstantiationer til separate kompileringsenheder, kræver de væsentlige tilpasninger af bygge-systemer og brugen af uhandy makroer. I modsætning hertil kræver CuTeDSL-kerner, der kompileres Just-In-Time (JIT), ikke komplekse bygningsregler og hæmmer ikke kompileringstiderne, da udviklingscyklussen kan fokuseres på et enkelt punkt i konfigurationsrummet, mens omhyggelige enhedstests senere kan dække hele kernens designrum.

Selvom JIT-kompilering flytter omkostninger fra udviklingstid til køretid, er overheaden mere end acceptabel. Når en model udrulles, aktiverer den kun få af punkterne fra hele konfigurationsrummet, hvilket medfører overheads i sekund-skala under opstart/opvarmning af inferensmotoren. I modsætning hertil kan præ-kompilering af alle konfigurationer af en kerne spænde over flere minutter. Derudover kan overheads yderligere afbødes ved at cache de JIT-kompilete artefakter.

De tunge kompileringsomkostninger ved CUDA forstærkes af vanskeligheden ved at spore oprindelsen af fejl, når skabelonbaseret dispatch og specialisering anvendes, hvilket er tilfældet med CuTe og CUTLASS. Med et DSL kan fejl peges på den linje kode, der genererede dem, hvilket letter fejlretningen væsentligt. Derudover er det meget enklere at tilknytte den genererede kode til den DSL-sætning, der frembragte den, via mellemliggende MLIR-stadier eller ved at stole på fejlretningsinformation knyttet til den genererede PTX. Takket være disse aspekter af DSL'en fandt vi os selv meget mere produktive, når vi udviklede kerner.

Stort konfigurationsrum

Selvom kerner deler mange ligheder på tværs af modeller, skal de stadig være udstyret med en række konfigurationsknapper, der kan have indflydelse på ydeevnen. Relevante parametre omfatter:

Skjult dimension/Hoveddimension/Antal attention-hoveder

Input-/Output-datatyper

Kvantiseringsskema

Compile-time-konstanter (bias, aktiveringsfunktioner osv.)

Udvid

Compile-time-specialisering er nyttig for både ydeevne og fejlretning. Især i en pipelinet processor, som det er tilfældet med GPU'er, er en ikke-eksisterende forgren bedre end en forgren, der altid eller aldrig tages. Derudover resulterer kompilering af en kerne til et specifikt sæt af konfigurationsknapper i færre PTX/SASS-instruktioner, som er lettere at analysere og feilsøge, når man bruger cuda-gdb under en undersøgelse.

For en given batchstørrelse og konfiguration kan den forventede trafik til DRAM og antallet af aritmetiske operationer estimeres. Derefter kan der træffes optimale beslutninger fra sag til sag for optimal gitterallokering, synkronisering og brug af visse hardwarefunktioner. CuTeDSL giver os mulighed for at definere og genbruge komponenter til dyb specialisering uden at pådrage os en uacceptabel kompilerings-overhead.

Forudfyldning og dekodningsspecialisering

Generelt skal ydeevnen af inferenskerner optimeres til to separate use cases:

Præfiks: lille batchstørrelse (få chunks, sekvenser), mange tokens pr. sekvens. Gennemstrømning skal maksimeres for både embeddings og input-tokens i LLM'er. Latens på ti til hundreder af millisekunder.

Dekodning: små til store batchstørrelser, en (ingen spekulativ dekodning) til et fåtal af tokens pr. sekvens (draft-target-spekulation, multi-token-forudsigelse). Latens tæt på kernestart-latens.

For GEMM og attention kører præfiks-fasen på et stort antal tokens, og ydeevnen er begrænset af enhedens beregningsgennemstrømning. Dekodning involverer den samme operation på små batchstørrelser, hvor eksekveringstiderne domineres af omkostningerne ved at indlæse aktiveringerne og vægtene fra DRAM ind i de streamende multiprocessorer og tensorkerner, begrænset af enhedens hukommelsesbåndbredde i stedet. Selvom batchstørrelsen eller det samlede antal tokens i en batch er en relativt klar indikator, er overgangstærsklen mellem præfiks og dekodning mere subtil, fordi hukommelsesbåndbredden også bestemmes af modellens skjulte dimension og aktiveringens datatype. Kerner opdeles derfor i præfiks- og dekoderimplementeringer, hvis algoritmiske forskelle eller gitter-/bloktildelinger ville favorisere en af dem.

Den parameter, der universelt er forskellig mellem præfiks- og dekoderkerner, er antallet af warps i en blok (antallet af tråde divideret med warp-størrelsen på 32 tråde pr. warp). Det er ønskeligt at gøre denne parameter til en compile-time-konstant, da det kan hjælpe med registerallokering. I CUDA gøres dette typisk ved at specificere antallet af tråde via __launch_bounds__ og levere antallet af warps som en skabelonparameter:

Antallet af registre til rådigheid for alle tråde i en blok er en hardwarebegrænsning: 64K på Hopper og Blackwell. Hvis antallet af tråde kendes på forhånd, ved kompilatoren præcis, hvor mange registre den kan tildele til hver tråd, og den kan få succes med at kompilere mere aggressivt optimerede kerner. Dette sker dog på bekostning af kompileringstid, da kernen nu skal genkompileres for forskellige warp-størrelser. Med CuTeDSL er dette ikke længere et problem, da vi ikke behøver at betale omkostningerne ved kompilering for hele konfigurationsrummet på forhånd. DSL'en forventer også, at blokstørrelsen er et konstant udtryk, hvilket implicit udfører den optimering, der er antydet af __latch_bounds__ i CUDA.

Udover at gøre antallet af warps til en konstant er det ønskeligt at variere det mellem præfiks og dekodning. For fuldt ud at udnytte en GPU's hukommelsesbåndbredde er det ikke tilstrækkeligt altid at tildele et højt antal warps pr. blok, da fuld båndbredde kun kan opnås ved at bruge så mange af de tilgængelige Streaming Multiprocessors (SM'er) som muligt. Følgelig vælger vi for et givet antal tokens og kendte datatyper og skjulte dimensioner først en ideel pr. tråd-vektordimension (8 for bfloat16, 4 for float32) for at sikre, at den genererede kode bruger 128-bit load/store-instruktionerne (256-bit på Blackwell). Når vi kender vektorstørrelsen, kan vi beregne det antal warps, der kræves af kernen. Vi prøver først at maksimere antallet af SM'er: Hvis der er færre warps end SM'er, vælges en enkelt warp pr. blok. Dette er typisk dekodervarianten af kernen. Hvis der er nok warps, øger vi gradvist antallet af warps til 8, 16 eller 32 og benchmarke en øvre grænse, hvorefter vi typisk ikke længere ser en forbedring i ydeevne.

Selvom Triton udførte nogle af disse optimeringer under kølerhjelmen, fandt vi ud af, at vi kan opnå lignende resultater ved omhyggeligt at bruge CuTe layout algebra uden at være begrænset af manglerne ved en high-level kompilator. Især når sub-warp horisontale reduktioner er involveret, fandt vi ud af, at CuTeDSL tilbød os den finkornede kontrol, der kræves for at generere optimal kode.

Gitter-synkronisering

I nogle situationer kan cooperative_groups::this_grid().sync() bruges til at synkronisere alle deltagende blokke for at distribuere arbejde på tværs af så mange blokke som muligt. Der gælder begrænsninger: Da en blok ikke kan af-skemalægges fra en SM, når den først begynder eksekvering, kan antallet af blokke for at garantere fremadrettet fremgang ikke overstige antallet af SM'er, der er til rådighed på enheden, for at sikre, at alle blokke på ethvert tidspunkt kan opdatere og afspørge de globale hukommelseslokationer, der er involveret i synkronisering. Denne begrænsning håndhæves af cudaLaunchCooperativeKernel. Omkostningerne ved en gitterbarriere er ca. 3us på Hopper plus omkostningerne ved at rydde alle tidligere hukommelsesoperationer via en implicit hukommelsesbarriere i visse situationer. Overordnet set kan dette sammenlignes med overheaden ved at opdele en kerne i to starter, hvilket fører til en ydeevne-følsom designdeslutning. På trods af det faktum, at Programmatic Dependent Launch (PDL) kan reducere opstarts-overheads fra og med Hopper, kan det betale sig for dekoderkerner at undgå overheaden.

Hvis en af arbejdsbyrderne kræver flere blokke end antallet af SM'er, bør kernen opdeles ved at stole på den implicitte barriere ved kernestarter. Dette er primært tilfældet for præfiks. Hvis det antal blokke, der er ideelt for alle arbejdsbyrder, passer ind i gitteret, foretrækkes en gitterbarriere, som det er tilfældet med dekoderkernerne, der opererer på et mindre antal tokens.

I CuTeDSL kan disse beslutninger tages i startkoden på værtssiden, og kernerne kan sammensættes og flettes med gitterbarriere efter behov, som illustreret af eksemplet nedenfor. Med en sådan hybrid strategi kan dekoderkerner optimeres til at nå minimale latenser, og præfiks-kerner kan optimeres til maksimal gennemstrømning. I modsætning til CUDA-baserede løsninger kan denne logik defineres mere kortfattet og med mindre kompilerings-overhead.

I tilfælde af vores MoE all-to-all-kerner, som beskrevet nedenfor, hjalp det os med at opnå væsentlige ydeevnegevinster gennem CuTeDSL at have fleksibiliteten til at bygge kerner med eller uden gitterbarriere og optimere dem ved kompileringstid.

CuTeDSL-kerner

QK-norm

Den første kerne, vi byggede i CuTeDSL, var en specialiseret RMS-normkerne, optimeret til små dimensioner. QK-norm, der bruges i self-attention-laget i de nyeste Qwen-modeller, anvendes pr. hoved på 64, 128 eller 256 elementer. I modsætning til almindelig RMS-norm, som opererer på skjulte dimensioner, der overstiger tusinder, hvilket kræver en horisontal reduktion gennem både delt hukommelse og registre, behøver QK-norm kun reduktion på tværs af warps med mulighed for at presse flere hoveder ind i en enkelt warp. Dette blev en use case, hvor vi kunne optimere aggressivt ved kompileringstid for at finde de bedste token-til-warp-tildelinger for alle hoveddimensioner og datatyper.

RMS-norm er en ret ligetil operation, selvom den kræver en horisontal sum for at beregne den kvadrede norm for en vektor:

Da attention-hoveder typisk er små og trådspecifikke operationer er stærkt vektoriserede, kan en enkelt warp beregne kvadratnormen. For at maksimere hukommelsesbåndbredden bør hver tråd behandle 128 bit input, hvilket svarer til 8 bfloat16-værdier. Følgelig er en halv-warp tilstrækkelig til at behandle et hoved, og den kvadrede norm kan bestemmes ved hjælp af en horisontal reduktion udført af shuffle_sync_bfly. Med udgangspunkt i hoveddimensionen og datatypen kan vi bestemme antallet af hoveder pr. warp og specificere argumenterne og antallet af reduceringsrunder på kompileringstidspunktet.

Vi kan derefter bygge en specialiseret kerne med relativt få linjer kode ved at stole på automatisk vektoriserede indlæsninger/lagre og bruge low-level arkitekturspecifikke operationer til warp-reduktionen. Kernen bestemmer, hvilket hoved (Q eller K) der skal arbejdes på, og identificerer de tilsvarende vægte, der skal anvendes. Den indlæser en hel vektor fra hovedet og summerer først lokalt. Derefter beregnes den kvadrede norm baseret på antallet af tråde pr. hoved. Til sidst bestemmes RMS-normen, der automatisk er vektoriseret, og skrives tilbage til hukommelsen. Modelspecifikke parametre såsom WEIGHT_BIAS og EPS er konstant-foldede.

Sammenlignet med en almindelig RMS-normkerne leverer denne specialiserede variant 2-3x bedre gennemstrømning og latens for både præfiks og dekodning.

Mixture-of-Experts Dispatch/Combine

MoE-dispatch og -combine shifter input-tokens til de eksperter, de routes til, og udfører en vægtet akkumulering, efter MoE GEMM-kerner udfører de relevante matrixmultiplikationer. Ydeevne er af særlig interesse, da forskellige produkter er afhængige af MoE-modeller i forskellige størrelser. Baseret på udrulningstyperne skelner vi mellem forskellige kerner, der skal optimeres:

TP=1: Lokal kerne til en enkelt GPU. Dispatch shifter tokens til de relevante eksperter, og combine akkumulerer dem. Bruges primært til Qwen3 30B-udrulninger på H100.

DP=1, TP ≠ 1: Tensor-parallel, singlenode-kerne. Inputtet replikeres, dispatch vælger tokens for den aktuelle rang, og combine akkumulerer dem på tværs af alle ranger. Akkumulering på tværs af enten NVLink eller NVLS. Bruges til store præfiks-generatorer.

DP ≠ 1, NVLink: Dataparallel kerne med valgfri tensor-parallelisme. Hver rang sender tokens til hver anden rang, der hoster specifikke eksperter. Routinginformation aggregeres, tokens overføres på tværs af NVIDIA NVLink. Bruges til dataparallelle H200- eller GB200-udrulninger.

DP ≠ 1, InfiniBand: Peer-to-Peer-dataparallel og tensor-parallel kerne, der opererer på tværs af både NVLink og NVIDIA Quantum InfiniBand. Bruges til multi-node Hopper-udrulninger.

Lokal

Den lokale kerne tæller antallet af tokens tildelt til hver ekspert baseret på routingtabellen og tildeler et unikt indeks til hvert token inden for listen af tokens, der tilhører hver ekspert. Efter routingtrinnet kopieres tokens ind i valgfrit polstrede buffere for at blive givet videre til en MoE GEMM-kerne. Ved hjælp af token-lokationer udfører combine-kernen derefter en akkumulering i registre og skriver de endelige tokens ind i en outputbuffer.

Dispatch-fasen varierer mellem præfiks og dekodning: For få tokens er der lidt routinginformation at læse ud, så vi behandler den i en enkelt blok og signalerer resten af gitteret, hvor individuelle blokke kopierer chunks af tokens. Token-tildeling bestemmes ved atomisk at inkrementere tællere i delt hukommelse. Da præfiks skal behandle mere routinginformation, opdeles routing på tværs af flere blokke i en separat kerne. Indekser tildeles først lokalt ved hjælp af tællere i delt hukommelse og aggregeres derefter på tværs af alle blokke i global hukommelse for at bestemme positionen af hvert token. Gitteret for denne kerne klemmes til antallet af SM'er for at kunne bruge gitter-vise barrierer. En efterfølgende kerne udfører kopieringerne, startet på tværs af et meget større antal blokke og warps for fuldt ud at udnytte den tilgængelige hukommelsesbåndbredde.

Combine-kernen er relativt enkel, da den udfører vektoriseret akkumulering med udgangspunkt i routinginformationen og de offsets, der er beregnet af dispatch. CuTeDSL viste sig værdifuld her, da beregningen kunne udtrykkes gennem enkle CuTe-tensorer, der kompilerer til effektiv kode for alle datatyper og skjulte dimensioner.

Singlenode

I en tensor-parallel udrulning replikeres routinginformationen og input-tokens. Den lokale dispatch-kerne kan genbruges fuldt ud med en lille ændring for at route og kun vælge de tokens, der er tildelt til de eksperter, der er placeret på den aktuelle rang. Med nogle parametre og compile-time-specialisering genbruges de lokale dispatch-kerner.

Da antallet af ranger normalt er lavere end det antal eksperter, et token routes til, kan det i det gennemsnitlige tilfælde antages, at hvert token har mindst én af sine replikaer routet til hver rang. I denne situation kan den lokale combine-kerne genbruges til at aggregere de lokale kopier af et token efterfulgt af en all-reduce for at summere på tværs af de deltagende ranger. I tilfælde af præfiks-kernen starter vi separate kerner, men dekodning fletter dog begge operationer ind i en enkelt start.

Til all-reduce bruger vi almindelig NVLink, hvis kun to ranger deltager, og skifter til NVLS, når 4 eller flere ranger er involveret, for en stigning på 10-20 % i gennemstrømning. Takket være DSL'en kan vi dele det meste af reduktionslogikken og peer-bufferudvekslingerne mellem MoE- og de almindelige all-reduce-kerner. Dette er en situation, hvor tæt kontrol over hardwaren er fordelagtig, da DSL-operationer over multimem.ld_reduce og multimem.st kræves for NVLS. Implementeringen af selve reduktionen følger i vid udstrækning all_reduce-eksemplet fra NCCL-tests.

Dataparallel

De dataparallelle kerner stoler på NVLink til at aggregere tokens til flere enheder. I dispatch-fasen, før tokens udveksles, routes peers først lokalt for at identificere måloffset på hver rang for hvert token, de sender. Derefter begynder de at kopiere tokens over i modtagerbuffere allokeret til hver afsender for at undgå konflikter. Pr. token release-store-skrivninger bruges til at signalere fuldførelse. På modtagersiden aggregeres routinginformation for at kunne kondensere alle tokens til en enkelt buffer til MoE GEMM. Modtagere afspørger det peer-skrevne flag med en acquire-load for at vente på, at skrivningerne fuldføres.

På combine-siden kopieres tokens over NVLink til kilderangerne med et pr. token-flag, der indikerer modtagelse. Combine-kernen svarer til de lokale og singlenode-varianter med små forskelle i adressering og noget ekstra synkroniseringslogik.

Vi byggede et bibliotek af PTX-indkapslinger for at eksponere release-acquire-indlæsninger og -lagre med omfang på både system- og enhedsniveau for at implementere denne kerne. Ved at forbedre granulariteten af synkroniseringsflag og allokere flere blokke til at sende og modtage tokens over NVLink opnåede vi 2x forbedringer i forhold til vores tidligere CUDA-kerner.

Peer-to-Peer

Vores Peer-to-Peer-kerner opnåede tidligere avanceret ydeevne på Hopper-GPU'er parret med ConnectX-7 InfiniBand-adaptere. Vi har siden portet combine-kernerne over til CuTeDSL for yderligere ydeevnegevinster. Da CUDA-versionerne allerede i høj grad stolede på compile-time-specialisering, viste de sig at være kompilerings-flaskehalsene i vores repository, hvilket tvang os til at udelade nogle parametre. Ved at migrere til det JIT-kompilete DSL specialiserede vi os i flere parametre og finjusterede gitterallokeringer, hvilket introducerede yderligere forskelle mellem præfiks og dekodning for omkring et boost på 3 % i latens.

Blackwell-understøttelse

Fra perspektivet af ROSE-kerner falder Blackwell-understøttelse i to kategorier: arkitekturaffektive og arkitekturspecifikke kerner. Arkitekturspecifikke kerner, der stoler på tensorkerner, kræver eksplicit arbejde for at tilpasse sig. Men takket være brugen af DSL'en tilpasser de arkitekturaffektive kerner sig trivielt.

En Blackwell-GPU har op til 160 SM'er i modsætning til Hoppers 132. For de fleste kerner er dette en compile-time-konstant, og logik i startkoden på værtssiden tager højde for det for at bestemme den optimale kerne (præfiks vs. dekodning) til at køre. Derudover introducerer Blackwell også 256-bit load-instruktioner, som let kan udsendes ved at øge vektorvidden (sidste dimension af CuTe-tensorer) for at forbedre udnyttelsen af hukommelsesbåndbredden.

Blackwell introducerer hardwareunderstøttelse af et betydeligt antal kvantiseringsskemaer (MXFP8, MXFP4, NVFP4). Understøttelse af disse blev primært implementeret ved at finjustere eksisterende kerner til FP8-blokvis skalering. JIT-understøttelse er afgørende her, da nye tilstande kan understøttes uden at eksplodere det kompileringsmæssige konfigurationsrum. Tidligere fandt vi Triton akavet til at udtrykke horisontale reduktioner på tværs af få elementer over halvdelen eller kvartdelen af en warp. Finkornet kontrol over reduktioner var afgørende for at opnå optimal ydeevne ved beregning af en absolut maksimumværdi på tværs af henholdsvis 16 eller 32 elementer for NVFP4 og MXFP8/MXFP4.

AI-assisteret kernegenerering

Som et relativt nyt sprog havde LLM'er ringe eksponering for CuTeDSL-kode under deres fortræning. Derudover var stilen i vores nuværende kode, med kerner optimeret til aktiveringer og normer, lidt forskellig fra de eksisterende CUTLASS- eller Quack-eksempler, som var stærkt optimerede til TMA og arkitekturspecifikke matrixmultiplikationsinstruktioner. At nå til et punkt, hvor LLM'er kunne hjælpe med kerneskrivning, var en gradvis proces.

Vi opbyggede manuelt et initialt korpus af kerner og hjælpefunktioner, der fulgte konsekvente mønstre, før vi kunne stole på agenter til kernekonvertering. Efterhånden som korpusset skred frem, kunne vi udføre grove oversættelser fra CUDA til CuTeDSL med betydeligt manuelt efterarbejde. Især fandt vi ud af, at modeller kæmpede med pointer-til-tensor-konverteringer og vektorisering og gik i stå i forsøget på at rette fejl i typer, når målet var at få en test til at kompilere. De har vist ringe forståelse for den gitterbaserede natur af kernerne og forveksler tråd- og gittersynkronisering. Mest frustrerende var der en fristelse til at hallucINere ikke-eksisterende oversættelser af CUDA-specifikke biblioteker i stedet for at bygge mod vores egne CuTeDSL-erstatninger. Over tid, efterhånden som korpusset af CuTeDSL-kerner voksede, fandt vi ud af, at LLM'er kan bygge det meste af skeletkoden og udfylde de tomme felter, hvilket gør os i stand til at fokusere på det overordnede design.

Konklusion og fremtidigt arbejde

Sammenlignet med CUDA finder vi, at CuTeDSL leverer meget bedre kompileringsydeevne, hvilket gør os i stand til at specialisere os i flere parametre og opnå ydeevneforbedringer. Selvom kerner er lidt mere ordrige på grund af manglen på visse abstraktioner sammenlignet med Triton, finder vi, at kompleksiteten er retfærdiggjort, da finkornet kontrol over hardwaren gør os i stand til at bygge hurtigere kerner.

På tidspunktet for skrivningen er ydeevnegabet mellem CUTLASS C++ og CuTeDSL-kerner minimalt, hvilket åbner op for spændende muligheder for at sammensætte og flette højtydende GEMM-kerner med inter-device-kommunikation, normalisering eller aktiveringer. Ved at forpligte os til denne DSL vil vi være i stand til at udtrække yderligere ydeevne fra vores enheder for at forbedre gennemstrømningen på tværs af vores arbejdsbyrder, fra embeddings til LLM'er.

Referencer

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/

Citering

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