CuTeDSL bij Perplexity

Bij Perplexity gebruiken we de in-house Runtime-Optimized Serving Engine (ROSE) om modellen te bedienen, variërend van embeddings tot LLM's met biljoenen parameters. Onze runtime, die achter alle API's draait, inclusief Sonar, Search en Embeddings, is gebouwd om onderzoekers in staat te stellen de meest geschikte modellen en architecturen te verkennen, terwijl de beste kandidaten snel in productie kunnen worden genomen. Om dit te bereiken, moeten we in staat zijn om inference voor state-of-the-art architecturen in een hoog tempo te implementeren, in te zetten en te optimaliseren. In deze blogpost onderzoeken we hoe CuTeDSL-kernels, geïntegreerd in onze inference engine, ons in staat stellen de vereiste GPU-kernels te bouwen om modellen naar topprestaties op NVIDIA GPU's te brengen.

Inference bij Perplexity AI

Aangepaste modellen, gebouwd door verschillende teams bij Perplexity, worden in-house gehost op NVIDIA Hopper en Blackwell GPU's met behulp van de ROSE-inference engine die is gebouwd door het AI Inference-team. ROSE is aanvankelijk gebouwd om op maat gemaakte Llama-modellen beter te bedienen voor zowel taaldecodering als classificatie, met een interface die compatibel is met de NVIDIA Triton Inference Server. Na verloop van tijd groeide het uit tot een volwaardige inference engine die substantiële LLM's bedient, evenals ranking-, classificatie-, scoring- en andere op transformatoren gebaseerde modellen voor een breed scala aan teams om de producten van Perplexity aan te sturen. Het ontwerpdoel van ROSE is eenvoudig te zijn: na verloop van tijd wordt er geïnvesteerd in het bereiken van topprestaties, maar de hoofdprioriteit is om ons aan te passen aan de steeds evoluerende architecturen en use-cases voor onze modellen.

Vergroten

Op hoofdlijnen bestaat ROSE uit een engine die een model aanpast aan een interface waar clients mee kunnen communiceren. De engines zijn zeer herbruikbaar, handelen het plannen en batchen van verzoeken af, en halen gedecodeerde tokens, embeddings of scores op uit de onderliggende modellen. De initialisatie van apparaten, inter-proces- en inter-node-communicatiekanalen en het laden van gewichten worden allemaal door hen afgehandeld. Voor LLM's is de engine verantwoordelijk voor batchen, chunking en sampling. Het is ook verantwoordelijk voor de toewijzing van KV-opslag, voor zowel volledige als lineaire aandacht, evenals prefix-matching. Embedding-engines zijn eenvoudiger, omdat ze alleen on-line batchen uitvoeren. De uitvoering is zodanig geoptimaliseerd dat de engine de volgende werkronde voorbereidt terwijl het model parallel op de accelerator draait.

ROSE stelt een rijke verzameling aangepaste lagen bloot voor modellering, waarbij GPU-kernels worden verpakt die zijn gebouwd met CuTeDSL, Triton, CUDA, CUTLASS en cuBLAS. Voor matrixvermenigvuldiging en aandachtsimplementaties vertrouwen we voornamelijk op NVIDIA-kernels van CUTLASS en cuBLAS. De rest van de pijplijn, inclusief embedding-, norm-, MoE- en activeringskernels, maakt deel uit van de kernenbibliotheek van ROSE. Om de hardware waarop we inference uitvoeren het beste te ondersteunen, is ROSE uitgerust met zijn eigen inter-device communicatie-primitieven voor MoE-routing, dispatch en combine. Hoewel de andere kernels relatief eenvoudig zijn, moeten ze grondig worden gespecialiseerd om topprestaties te behalen op alle verschillende verborgen dimensies en apparaten waarop ze draaien. Bovendien vertrouwen we op compile-time specialisatie om rekening te houden met alle knoppen en aanpassingen van verschillende modellen (weight bias versus geen weight bias voor RMS-norm, activeringsfuncties, enz.). Bijgevolg hebben we besloten om DSL's te gebruiken om ze te implementeren: we hebben ROSE aanvankelijk gebouwd met Triton. We hebben sindsdien echter CuTeDSL geadopteerd als de primaire omgeving voor GPU-programmering.

Waarom CuTeDSL?

Vanuit wiskundig perspectief zijn inference-kernels meestal ongecompliceerd. Bij het ontwikkelen van een kernel ligt de complexiteit daarom over het algemeen niet in het uitdrukken van welke berekening moet worden uitgevoerd, maar in hoe de hardware dit precies moet doen. Om de beste kernels te bouwen, hebben we een taal nodig die de juiste abstracties biedt over de onderliggende hardware. De juiste taal stelt ons in staat om de high-level operaties te specificeren om de berekening uit te voeren, waarbij ze agressief worden gespecialiseerd en geoptimaliseerd over alle punten in een brede configuratieruimte.

CuTeDSL is een Python-gebaseerde DSL voor NVIDIA GPU's, gebouwd op basis van CuTe-lay-outalgebra en MLIR. Het compileert just-in-time naar sterk geoptimaliseerde PTX, terwijl de controle over hardwareprimitieven op laag niveau behouden blijft. Het beoogt dezelfde functies en prestaties te bieden als CUTLASS C++, zonder de dure compilatietijden die de ontwikkelingssnelheid hinderen. Van de talen die we hebben onderzocht of waarmee we hebben gewerkt (inclusief CUDA, cuTile, TVM IR en Triton), beschouwen we dit als de beste pasvorm voor de ontwikkeling van onze kernels.

Over het algemeen, hoewel niet zo high-level en beknopt als Triton of cuTile, vonden we dat de CuTe-lay-outalgebra en de composities daarvan (blootgesteld via cute.Tensor) redelijke vervangers waren. Bovendien biedt CuTeDSL onbeperkte toegang tot hardwareprimitieven, wat ons het niveau van controle biedt dat we nodig hebben om in alle randgevallen topprestaties te behalen. Met CUDA waren de mogelijkheden voor compile-time specialisatie beperkt, omdat elke extra parameter onredelijke kosten voor compilatietijden of complexiteit toevoegde aan het buildsysteem. In tegenstelling hiermee kunnen we met een Just-In-Time gecompileerde DSL agressiever specialiseren zonder de iteratiesnelheid op te offeren. In tegenstelling tot TVM IR, dat CUDA als een tussenstap gebruikte, vonden we de op MLIR gebaseerde pijplijn die verlaagde naar PTX robuuster en gemakkelijker te debuggen dan een pijplijn die CUDA als tussenstap uitzond. Bovendien ontdekten we, met het gebruik van aangepaste operaties om PTX-instructies in te pakken, dat we alle functies van de onderliggende hardware volledig konden benutten.

Ontwikkelaarservaring

Het grootste deel van de compilatiekosten van CUDA-kernels wordt doorgaans besteed aan de fase van sjabloonuitbreiding van kernels, aangezien optimalisatie en compilatie vanaf de PTX-fase kunnen worden voltooid in een kwestie van seconden. Met veel compile-time parameters kan deze compilatiefase exploderen tot meerdere minuten voor een niet-triviale kernel. Dit knelpunt wordt geaccentueerd door het feit dat mogelijkheden voor parallelle compilatie pas later ontstaan, op IR- en PTX-niveau, waarbij de afhandeling van C++-sjablonen sequentieel plaatsvindt. Hoewel er trucs bestaan om dit aan te pakken door sjablooninstantiaties naar afzonderlijke compilatie-eenheden te verplaatsen, vereisen ze aanzienlijke aanpassingen aan buildsytemen en het gebruik van onhandige macro's. Daarentegen vereisen CuTeDSL-kernels die Just-In-Time (JIT) worden gecompileerd geen complexe buildregels en hinderen ze de compilatietijden niet, aangezien de ontwikkelingscyclus kan worden gefocust op een enkel punt in de configuratieruimte, terwijl zorgvuldige unit-tests later de gehele ontwerpruimte van de kernel kunnen bestrijken.

Hoewel JIT-compilatie kosten verplaatst van ontwikkelingstijd naar runtime, is de overhead meer dan acceptabel. Wanneer een model wordt ingezet, activeert het slechts enkele punten uit de volledige configuratieruimte, wat seconden-schaal overhead veroorzaakt tijdens het opstarten/opwarmen van de inference engine. Daarentegen zou het vooraf compileren van alle configuraties van een kernel meerdere minuten kunnen duren. Bovendien kunnen overheadkosten verder worden beperkt door de JIT-gecompileerde artefacten in de cache op te slaan.

De hoge compilatiekosten van CUDA worden geaccentueerd door de moeilijkheid om de oorsprong van fouten te achterhalen wanneer sjabloongebaseerde dispatch en specialisatie worden gebruikt, zoals het geval is bij CuTe en CUTLASS. Met een DSL kunnen fouten worden aangewezen op de regel code die deze heeft gegenereerd, wat het debuggen aanzienlijk vergemakkelijkt. Bovendien is het veel eenvoudiger om de gegenereerde code toe te wijzen aan de DSL-verklaring die deze heeft gegenereerd via tussenliggende MLIR-fasen of door te vertrouwen op debug-informatie die aan de gegenereerde PTX is gekoppeld. Dankzij deze aspecten van de DSL vonden we onszelf veel productiever bij het ontwikkelen van kernels.

Grote configuratieruimte

Hoewel kernels veel overeenkomsten delen tussen modellen, moeten ze nog steeds worden uitgerust met een reeks configuratieknoppen die van invloed kunnen zijn op de prestaties. Relevante parameters zijn onder meer:

Verborgen dimensie/hoofddimensie/aantal aandachtshoofden

Invoer/uitvoer-datatypen

Kwantiseringsschema

Compile-time constanten (biases, activeringsfuncties, enz.)

Vergroten

Compile-time specialisatie is nuttig voor zowel prestaties als debuggen. Vooral in een gepipelineerde processor, zoals het geval is bij GPU's, is een niet-bestaande tak beter dan een altijd genomen of nooit genomen tak. Bovendien resulteert het compileren van een kernel voor een specifieke set configuratieknoppen in minder PTX/SASS-instructies, die gemakkelijker te analyseren en te debuggen zijn bij het gebruik van cuda-gdb tijdens een onderzoek.

Voor een gegeven batchgrootte en configuratie kunnen het verwachte verkeer naar DRAM en het aantal rekenkundige operaties worden geschat. Vervolgens kunnen er per geval optimale beslissingen worden genomen voor optimale grid-toewijzing, synchronisatie en het gebruik van bepaalde hardwarefuncties. CuTeDSL stelt ons in staat om componenten te definiëren en te hergebruiken voor diepe specialisatie zonder onaanvaardbare compile-time overhead te veroorzaken.

Specialisatie in prefill en decode

Over het algemeen moeten de prestaties van inference-kernels worden geoptimaliseerd voor twee afzonderlijke use-cases:

Prefill: kleine batchgrootte (weinig chunks, sequenties), veel tokens per sequentie. Doorvoer moet worden gemaximaliseerd voor zowel embeddings als de invoertokens van LLM's. Latency in de tientallen tot honderden milliseconden.

Decode: batchgroottes van klein tot groot, één (geen speculatieve decodering) tot enkele tokens per sequentie (draft-target speculatie, multi-token voorspelling). Latency dicht bij de opstartlatentie van de kernel.

Voor GEMM en aandacht draait de prefill-fase op een groot aantal tokens en worden prestaties begrensd door de rekenkracht van het apparaat. Decoding omvat dezelfde operatie op kleine batchgroottes, waarbij uitvoeringstijden worden gedomineerd door de kosten van het laden van de activaties en gewichten van DRAM naar de streaming multiprocessors en tensor-cores, beperkt door de geheugenbandbreedte van het apparaat in plaats daarvan. Hoewel de batchgrootte of het totale aantal tokens in een batch een relatief duidelijke indicator is, is de drempelwaarde voor de overschakeling tussen prefill en decode subtieler omdat de geheugenbandbreedte ook wordt bepaald door de verborgen dimensie van het model en het datatype van de activering. Kernels worden daarom opgesplitst in prefill- en decode-implementaties als algoritmische verschillen of grid/blok-toewijzingen een van beide zouden bevoordelen.

De parameter die universeel verschilt tussen prefill- en decode-kernels is het aantal warps in een blok (het aantal draden, gedeeld door de warp-grootte van 32 draden per warp). Het is wenselijk om van deze parameter een compile-time constante te maken, omdat dit kan helpen bij registertoewijzing. In CUDA wordt dit meestal gedaan door het aantal draden te specificeren via __launch_bounds__ en het aantal warps als sjabloonparameter op te geven:

Het aantal registers dat beschikbaar is voor alle draden in een blok is een hardwarebeperking, 64K op Hopper en Blackwell. Als het aantal draden vooraf bekend is, weet de compiler precies hoeveel registers hij aan elke draad kan toewijzen en kan hij erin slagen agressiever geoptimaliseerde kernels te compileren. Dit gaat echter gepaard met compilatiekosten, omdat de kernel nu opnieuw moet worden gecompileerd voor verschillende warp-groottes. Met CuTeDSL is dit geen probleem meer, aangezien we de kosten van compilatie voor de gehele configuratieruimte niet vooraf hoeven te betalen. De DSL verwacht ook dat de blokgrootte een constante expressie is, waarbij impliciet de optimalisatie wordt uitgevoerd die wordt geïmpliceerd door __latch_bounds__ in CUDA.

Naast het constant maken van het aantal warps, is het wenselijk om te variëren tussen prefill en decode. Om de geheugenbandbreedte van een GPU volledig te benutten, is het niet voldoende om altijd een hoog aantal warps per blok toe te wijzen, aangezien volledige bandbreedte alleen kan worden bereikt door zoveel mogelijk van de beschikbare Streaming Multiprocessors (SM's) te gebruiken. Daarom kiezen we voor een bepaald aantal tokens en bekende datatypen en verborgen dimensies eerst een ideale per-draad vectordimensie (8 voor bfloat16, 4 voor float32) om ervoor te zorgen dat de gegenereerde code de 128-bits load/store-instructies (256-bits op Blackwell) gebruikt. Met de vectorgrootte bekend, kunnen we het aantal warps berekenen dat door de kernel vereist is. We proberen eerst het aantal SM's te maximaliseren: als er minder warps zijn dan SM's, wordt een enkele warp per blok gekozen. Dit is doorgaans de decode-variant van de kernel. Als er genoeg warps zijn, verhogen we het aantal warps geleidelijk naar 8, 16 of 32, waarbij we een bovengrens benchmarken waarna we doorgaans geen verbetering in prestaties meer zien.

Hoewel Triton enkele van deze optimalisaties onder de motorkap uitvoerde, vonden we dat we vergelijkbare resultaten konden bereiken door zorgvuldig gebruik te maken van CuTe-lay-outalgebra, zonder beperkt te worden door de tekortkomingen van een high-level compiler. Vooral wanneer sub-warp horizontale reducties betrokken zijn, vonden we dat CuTeDSL ons de fijnmazige controle bood die nodig was om optimale code te genereren.

Grid-synchronisatie

In sommige situaties, om werk over zoveel mogelijk blokken te verdelen, kan cooperative_groups::this_grid().sync() worden gebruikt om alle deelnemende blokken te synchroniseren. Er gelden beperkingen: aangezien een blok niet van een SM kan worden gede-scheduleerd zodra de uitvoering begint, mag het aantal blokken om vooruitgang te garanderen niet hoger zijn dan het aantal SM's dat op het apparaat beschikbaar is, om ervoor te zorgen dat alle blokken op elk punt de globale geheugenlocaties kunnen bijwerken en pollen die bij synchronisatie betrokken zijn. Deze beperking wordt afgedwongen door cudaLaunchCooperativeKernel. De kosten van een grid-barrière zijn ongeveer 3us op Hopper, plus de kosten van het flushen van alle eerdere geheugenoperaties via een impliciete geheugenbarrière in sommige situaties. Over het algemeen is dit vergelijkbaar met de overhead van het splitsen van een kernel in twee lanceringen, wat leidt tot een prestatiegevoelige ontwerpbeslissing. Ondanks het feit dat Programmatic Dependent Launch (PDL) lanceringsoverhead vanaf Hopper kan verminderen, is het voor decode-kernels de moeite waard om de overhead te vermijden.

Als een van de workloads meer blokken vereist dan het aantal SM's, moet de kernel worden gesplitst, waarbij wordt vertrouwd op de impliciete barrière van kernel-lanceringen. Dit is voornamelijk het geval voor prefill. Als het aantal blokken dat ideaal is voor alle workloads in het grid past, heeft een grid-barrière de voorkeur, zoals het geval is bij de decode-kernels die op een kleiner aantal tokens werken.

In CuTeDSL kunnen deze beslissingen worden genomen in de host-side lanceringscode en kunnen de kernels naar behoefte worden samengesteld en afgewisseld met grid-barrières, zoals geïllustreerd door het onderstaande voorbeeld. Met een dergelijke hybride strategie kunnen decode-kernels worden geoptimaliseerd om minimale latenties te bereiken en kunnen prefill-kernels worden geoptimaliseerd voor maximale doorvoer. In tegenstelling tot CUDA-gebaseerde oplossingen kan deze logica beknopter en met minder compile-time overhead worden gedefinieerd.

In het geval van onze MoE all-to-all kernels, zoals hieronder beschreven, hielp het hebben van de flexibiliteit om kernels met of zonder grid-barrières te bouwen en ze bij compilatie te optimaliseren ons aanzienlijke prestatiewinsten te behalen via CuTeDSL.

CuTeDSL-kernels

QK Norm

De eerste kernel die we bouwden in CuTeDSL was een gespecialiseerde RMS-norm-kernel, geoptimaliseerd voor kleine dimensies. QK-norm, gebruikt in de self-attention-laag van de recente Qwen-modellen, wordt toegepast op basis van per-hoofd, op 64, 128 of 256 elementen. In tegenstelling tot reguliere RMS-norm, die werkt op verborgen dimensies van meer dan duizenden, waarbij een horizontale reductie via zowel gedeeld geheugen als registers vereist is, heeft QK-norm alleen reductie via warps nodig, met de mogelijkheid om meerdere hoofden in een enkele warp te persen. Dit werd een use-case waar we agressief konden optimaliseren tijdens het compileren om de beste token-to-warp-toewijzingen voor alle hoofddimensies en datatypen te vinden.

RMS-norm is een vrij eenvoudige operatie, hoewel het een horizontale som vereist om de gekwadrateerde norm van een vector te berekenen:

Aangezien aandachtshoofden doorgaans klein zijn en per-draad operaties sterk gevectoriseerd zijn, kan een enkele warp de kwadraatnorm berekenen. Om de geheugenbandbreedte te maximaliseren, moet elke draad 128 bits aan invoer verwerken, wat zich vertaalt naar 8 bfloat16 waarden. Bijgevolg is een halve warp voldoende om een hoofd te verwerken, en kan de gekwadrateerde norm worden bepaald met behulp van een horizontale reductie uitgevoerd door shuffle_sync_bfly. Vanaf de hoofddimensie en het datatype kunnen we het aantal hoofden per warp bepalen en de argumenten en het aantal reductierondes op het moment van compilatie specificeren.

We kunnen dan een gespecialiseerde kernel bouwen met relatief weinig regels code, vertrouwend op automatisch gevectoriseerde loads/stores en het gebruik van architectuurspecifieke operaties op lager niveau voor de warp-reductie. De kernel bepaalt aan welk hoofd (Q of K) moet worden gewerkt en identificeert de bijbehorende gewichten die moeten worden toegepast. Het laadt een volledige vector van het hoofd en telt eerst lokaal op. Vervolgens wordt op basis van het aantal draden per hoofd de kwadraatnorm berekend. Ten slotte wordt de RMS-norm, automatisch gevectoriseerd, bepaald en teruggeschreven naar het geheugen. Modelspecifieke parameters, zoals WEIGHT_BIAS en EPS, worden constant-folded.

Vergeleken met een reguliere RMS-norm-kernel levert deze gespecialiseerde variant 2-3x betere doorvoer en latentie voor zowel prefill als decode.

Mixture-of-Experts Dispatch/Combine

MoE-dispatch en combine schuiven invoertokens naar de experts waarnaar ze worden gerouteerd, waarbij een gewogen accumulatie wordt uitgevoerd nadat MoE GEMM-kernels de relevante matrixvermenigvuldigingen hebben uitgevoerd. Prestaties zijn van bijzonder belang, aangezien verschillende producten vertrouwen op MoE-modellen van verschillende groottes. Op basis van de inzettypen onderscheiden we verschillende kernels om te optimaliseren:

TP=1: Lokale kernel voor een enkele GPU. Dispatch schuift tokens naar de relevante experts en combine accumuleert ze. Voornamelijk gebruikt voor Qwen3 30B-implementaties op H100.

DP=1, TP ≠ 1: Tensor-parallelle, single-node kernel. De invoer wordt gerepliceerd, dispatch selecteert de tokens voor de huidige rangorde en combine accumuleert ze over alle rangordes heen. Accumulatie via NVLink of NVLS. Gebruikt voor grote prefillers.

DP ≠ 1, NVLink: Data-parallelle kernel met optionele tensor-parallellisme. Elke rangorde stuurt tokens naar elke andere rangorde die specifieke experts host. Routinginformatie wordt geaggregeerd, tokens worden overgedragen via NVIDIA NVLink. Gebruikt voor data-parallelle H200- of GB200-implementaties.

DP ≠ 1, InfiniBand: Peer-to-Peer data-parallelle en tensor-parallelle kernel die werkt via zowel NVLink als NVIDIA Quantum InfiniBand. Gebruikt voor multi-node Hopper-implementaties.

Lokaal

De lokale kernel telt het aantal tokens dat aan elke expert is toegewezen op basis van de routeringstabel en wijst een unieke index toe aan elk token binnen de lijst met tokens die bij elke expert hoort. Na de routeringsstap worden de tokens gekopieerd naar optioneel opgevulde buffers die moeten worden doorgegeven aan een MoE GEMM-kernel. Met behulp van de tokenlocaties voert de combine-kernel vervolgens een accumulatie uit in registers, waarbij de uiteindelijke tokens in een uitvoerbuffer worden geschreven.

De dispatch-fase varieert tussen prefill en decode: voor weinig tokens is er weinig routinginformatie om uit te lezen, dus verwerken we deze in een enkel blok en signaleren we de rest van het grid, waarbij individuele blokken stukjes tokens kopiëren. Token-toewijzing wordt bepaald door tellers in gedeeld geheugen atomair te verhogen. Omdat prefill meer routinginformatie moet verwerken, wordt routing opgesplitst over meerdere blokken in een aparte kernel. Indices worden eerst lokaal toegewezen met behulp van tellers in gedeeld geheugen en worden vervolgens geaggregeerd over alle blokken in het globale geheugen om de positie van elk token te bepalen. Het grid van deze kernel is vastgeklemd op het aantal SM's om grid-wijze barrières te kunnen gebruiken. Een daaropvolgende kernel voert de kopieën uit, gelanceerd over een veel groter aantal blokken en warps om de beschikbare geheugenbandbreedte volledig te benutten.

De combine-kernel is relatief eenvoudig, aangezien deze gevectoriseerde accumulatie uitvoert, vertrouwend op de routinginformatie en op de offsets berekend door dispatch. CuTeDSL bleek hier waardevol, aangezien de berekening kon worden uitgedrukt via eenvoudige CuTe-tensoren die compileren naar efficiënte code voor alle datatypen en verborgen dimensies.

Single-node

In een tensor-parallelle inzet worden de routinginformatie en invoertokens gerepliceerd. De lokale dispatch-kernel kan volledig worden hergebruikt met een kleine wijziging om alleen de tokens te routeren en te selecteren die zijn toegewezen aan de experts die zich op de huidige rangorde bevinden. Met enkele parameters en compile-time specialisatie worden de lokale dispatch-kernels hergebruikt.

Aangezien het aantal rangordes meestal lager is dan het aantal experts waarnaar een token wordt gerouteerd, kan in het gemiddelde geval worden aangenomen dat elk token ten minste een van zijn replica's naar elke rangorde heeft gerouteerd. In deze situatie kan de lokale combine-kernel worden hergebruikt om de lokale kopieën van een token te aggregeren, gevolgd door een all-reduce om de som te bepalen over de deelnemende rangordes heen. In het geval van de prefill-kernel lanceren we afzonderlijke kernels, maar decode fuseert beide operaties in een enkele lancering.

Voor all-reduce gebruiken we gewone NVLink als slechts twee rangordes deelnemen, en schakelen we over naar NVLS wanneer 4 of meer rangordes betrokken zijn voor een toename van 10-20% in doorvoer. Dankzij de DSL kunnen we het grootste deel van de reductielogica en de peer-buffer-uitwisselingen tussen de MoE en de reguliere all-reduce-kernels delen. Dit is een situatie waarin nauwe controle over de hardware gunstig is, aangezien DSL-operaties via multimem.ld_reduce en multimem.st vereist zijn voor NVLS. De implementatie van de reductie zelf volgt grotendeels het all_reduce-voorbeeld uit NCCL-tests.

Data-parallel

De Data-parallelle kernels vertrouwen op NVLink om tokens voor meerdere apparaten te aggregeren. In de dispatch-fase, voordat tokens worden uitgewisseld, routeren peers eerst lokaal om de doel-offset op elke rangorde te identificeren voor elk token dat ze verzenden. Vervolgens beginnen ze de tokens te kopiëren naar ontvangstbuffers die voor elke afzender zijn toegewezen om conflicten te voorkomen. Per-token release-store-writes worden gebruikt om voltooiing te signaleren. Aan de kant van de ontvanger wordt routeringsinformatie geaggregeerd om alle tokens in een enkele buffer te kunnen condenseren voor MoE GEMM. Ontvangers pollen de door de peer geschreven vlag met een acquire-load om te wachten tot de schrijfbewerkingen zijn voltooid.

Aan de combine-kant worden tokens gekopieerd via NVLink naar de bronrangordes, met een per-token vlag die de ontvangst aangeeft. De combine-kernel is vergelijkbaar met de lokale en single-node varianten, met kleine verschillen in adressering en wat extra synchronisatielogica.

We hebben een bibliotheek van PTX-wrappers gebouwd om release-acquire-loads en stores bloot te stellen die zowel op systeem- als apparaatniveau zijn gescopeerd om deze kernel te implementeren. Door de granulariteit van synchronisatievlaggen te verbeteren en meer blokken toe te wijzen om tokens via NVLink te verzenden en te ontvangen, hebben we 2x verbeteringen behaald ten opzichte van onze eerdere CUDA-kernels.

Peer-to-Peer

Onze Peer-to-Peer-kernels behaalden voorheen state-of-the-art prestaties op Hopper GPU's gekoppeld aan ConnectX-7 InfiniBand-adapters. We hebben sindsdien de combine-kernels geport naar CuTeDSL voor verdere prestatiewinst. Aangezien de CUDA-versies al zwaar vertrouwden op compile-time specialisatie, bleken ze de knelpunten van de compilatietijd in onze repository te zijn, waardoor we sommige parameters moesten weglaten. Door te migreren naar de JIT-gecompileerde DSL, specialiseerden we ons voor meer parameters en pasten we grid-toewijzingen aan, wat verdere verschillen tussen prefill en decode introduceerde voor ongeveer een 3% boost aan latentie.

Blackwell-ondersteuning

Vanuit het perspectief van ROSE-kernels valt Blackwell-ondersteuning in twee categorieën: architectuur-agnostische en architectuurspecifieke kernels. Architectuurspecifieke kernels, die vertrouwen op Tensor-cores, vereisen expliciet werk om zich aan te passen. Dankzij het gebruik van de DSL passen de architectuur-agnostische kernels zich echter triviaal aan.

Een Blackwell GPU heeft tot 160 SM's, in tegenstelling tot de 132 van Hopper. Voor de meeste kernels is dit een compile-time constante en logica in de host-side lanceringscode houdt er rekening mee bij het bepalen van de optimale (prefill versus decode) kernel om te draaien. Bovendien introduceert Blackwell ook 256-bits load-instructies, die gemakkelijk kunnen worden uitgezonden door de vectorbreedte (laatste dimensie van CuTe-tensoren) te vergroten om het gebruik van de geheugenbandbreedte te verbeteren.

Blackwell introduceert hardware-ondersteuning voor een aanzienlijk aantal kwantiseringsschema's (MXFP8, MXFP4, NVFP4). Ondersteuning hiervoor werd voornamelijk geïmplementeerd door bestaande kernels aan te passen voor FP8 blok-wijze schaling. JIT-ondersteuning is hier cruciaal, aangezien nieuwe modi kunnen worden ondersteund zonder de compile-time configuratieruimte te laten exploderen. In het verleden vonden we Triton onhandig bij het uitdrukken van horizontale reducties over enkele elementen, over de helft of een kwart van een warp. Fijnmazigere controle over reducties was cruciaal om optimale prestaties te behalen bij het berekenen van een absolute maximumwaarde over respectievelijk 16 of 32 elementen voor NVFP4 en MXFP8/MXFP4.

AI-ondersteunde kernelgeneratie

Als relatief nieuwe taal hadden LLM's weinig blootstelling aan CuTeDSL-code tijdens hun pre-training. Bovendien was de stijl van onze huidige code, met kernels geoptimaliseerd voor activaties en normen, enigszins anders dan de bestaande CUTLASS- of Quack-voorbeelden, die zwaar waren geoptimaliseerd voor TMA en architectuurspecifieke matrixvermenigvuldigingsinstructies. Het bereiken van een punt waarop LLM's konden assisteren bij het schrijven van kernels was een geleidelijk proces.

We hebben handmatig een initiële corpus van kernels en hulpprogramma's opgebouwd, volgens consistente patronen, voordat we konden vertrouwen op agents voor kernelconversie. Naarmate het corpus vorderde, konden we ruwe vertalingen uitvoeren van CUDA naar CuTeDSL, met aanzienlijk handmatig herwerk. In het bijzonder ontdekten we dat modellen worstelden met pointer-to-tensor conversies en vectorisatie, waarbij ze vastliepen bij het proberen typefouten te corrigeren om een test te laten compileren. Ze toonden weinig begrip van de grid-gebaseerde aard van de kernels, waarbij ze draad- en grid-synchronisatie verwarden. Het meest frustrerend was de verleiding om niet-bestaande vertalingen van CUDA-specifieke bibliotheken te hallucineren in plaats van te bouwen tegen onze eigen CuTeDSL-vervangingen. Naarmate het corpus van CuTeDSL-kernels groeide, ontdekten we na verloop van tijd dat LLM's het grootste deel van de skeletcode kunnen bouwen en de lege plekken kunnen invullen, waardoor we ons kunnen concentreren op het high-level ontwerp.

Conclusie en toekomstig werk

Vergeleken met CUDA vinden we dat CuTeDSL veel betere compile-time prestaties levert, waardoor we ons kunnen specialiseren voor meer parameters en prestatieverbeteringen kunnen behalen. Hoewel kernels iets uitgebreider zijn vanwege het ontbreken van bepaalde abstracties in vergelijking met Triton, vinden we dat de complexiteit gerechtvaardigd is omdat fijnmazigere controle over de hardware ons in staat stelt snellere kernels te bouwen.

Op het moment van schrijven is het prestatieverschil tussen CUTLASS C++ en CuTeDSL-kernels minimaal, wat opwindende mogelijkheden opent om high-performance GEMM-kernels samen te stellen en te fuseren met inter-device communicatie, normalisatie of activaties. Door ons te verbinden aan deze DSL, zullen we in staat zijn om verdere prestaties uit onze apparaten te halen om de doorvoer van onze workloads te verbeteren, van embeddings tot LLM's.

Referenties

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/

Citaat

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