CuTeDSL at Perplexity

På Perplexity använder vi den egna körtidsoptimerade serveringsmotorn Runtime-Optimized Serving Engine (ROSE) för att köra modeller som sträcker sig från inbäddningar till språkmodeller (LLM) med biljoner parametrar. Vår körtid ligger bakom alla API:er, inklusive Sonar, Search och Embeddings, och har byggts för att ge forskare möjlighet att utforska de lämpligaste modellerna och arkitekturerna, samtidigt som de snabbt kan driftsätta de bästa kandidaterna i produktion. För att uppnå detta måste vi snabbt kunna implementera, driftsätta och optimera inferens för banbrytande arkitekturer. I det här blogginlägget utforskar vi hur CuTeDSL-kärnor som integrerats i vår inferensmotor gör det möjligt för oss att bygga de GPU-kärnor som krävs för att få modellerna att nå topprestanda på NVIDIA-GPU:er.

Anpassade modeller, byggda av olika team på Perplexity, hostas internt på NVIDIA Hopper- och Blackwell-GPU:er med hjälp av ROSE-inferensmotorn som har byggts av teamet för AI-inferens. ROSE byggdes ursprungligen för att bättre betjäna anpassade Llama-modeller för både språkavkodning och klassificering, med ett gränssnitt som är kompatibelt med NVIDIA Triton Inference Server. Med tiden växte den till en fullfjädrad inferensmotor som hanterar omfattande språkmodeller samt rangordnings-, klassificerings-, poängsättnings- och andra transformerbaserade modeller för ett brett spektrum av team för att driva Perplexitys produkter. Designmålet för ROSE är att vara enkelt: med tiden investeras ansträngningar för att nå topprestanda, men den främsta prioriteten är att anpassa sig till de ständigt föränderliga arkitekturerna och användningsfallen för våra modeller.

Kärnans konfigurationsutrymme över datatyper, huvuddimensioner, kompileringstidskonstanter och startinställningar.

Övergripande består ROSE av en motor som anpassar en modell till ett gränssnitt som klienter kan interagera med. Motorerna är i hög grad återanvändbara och hanterar schemaläggning och paketering av förfrågningar samt hämtning av avkodade tokens, inbäddningar eller poäng från de underliggande modellerna. De hanterar all initiering av enheter, kommunikationskanaler mellan processer och noder samt laddning av vikter. För språkmodeller ansvarar motorn för paketering, uppdelning i sjok och samplign. Den ansvarar även för allokering av KV-lagring, för både full och linjär attention, samt matchning av prefix. Inbäddningsmotorer är enklare eftersom de endast utför online-paketering. Exekveringen är optimerad så att motorn förbereder nästa runda av arbete medan modellen körs parallellt på acceleratorn.

ROSE exponerar en rik samling anpassade lager för modellering, som kapslar in GPU-kärnor byggda med CuTeDSL, Triton, CUDA, CUTLASS och cuBLAS. För matrismultiplikation och attention-implementeringar förlitar vi oss i första hand på NVIDIA-kärnor från CUTLASS och cuBLAS. Resten av pipelinen, inklusive inbäddning, normering, MoE- och aktiveringskärnor, ingår i ROSEs kärnbibliotek. För att på bästa sätt stödja den maskinvara vi kör inferens på är ROSE utrustad med egna interdevice-kommunikationsprimitiver för MoE-routing, -dispatch och -kombinering. Även om de andra kärnorna är relativt enkla måste de specialiseras noggrant för att uppnå topprestanda för alla de olika dolda dimensioner och enheter de körs på. Dessutom förlitar vi oss på kompileringsspecialisering för att hantera alla reglage och finjusteringar hos olika modeller (viktbias jämfört med ingen viktbias för RMS-normering, aktiveringsfunktioner med mera). Följaktligen beslutade vi att använda DSL:er för att implementera dem: vi byggde ursprungligen ROSE med Triton. Vi har dock sedan dess gått över till CuTeDSL som primär miljö för GPU-programmering.

Ur ett matematiskt perspektiv är inferenskärnor vanligtvis okomplicerade. Följaktligen ligger komplexiteten vid utveckling av en kärna i anslutning till detta i allmänhet inte i att uttrycka vilken beräkning som ska utföras, utan i exakt hur maskinvaran bör göra det. För att bygga de bästa kärnorna behöver vi ett språk som erbjuder rätt abstraktioner för den underliggande maskinvaran. Rätt språk gör det möjligt för oss att specificera de övergripande åtgärder som krävs för att utföra beräkningen, och att aggressivt specialisera och optimera dem över alla punkter i ett brett konfigurationsutrymme.

CuTeDSL är ett Python-baserat DSL för NVIDIA-GPU:er som byggts ovanpå CuTe-layoutalgebra och MLIR. Det kompileras just-in-time till starkt optimerad PTX, samtidigt som kontrollen över maskinvaruprimitiver på låg nivå behålls. Målet är att erbjuda samma funktioner och prestanda som CUTLASS C++, utan de tidsödande kompileringstider som hämmar utvecklingstakten. Av de språk vi har utvärderat eller arbetat med (inklusive CUDA, cuTile, TVM IR och Triton) anser vi att det passar bäst för utvecklingen av våra kärnor.

Sammantaget fann vi att CuTe Layout Algebra och dess sammansättningar (exponerade via cute.Tensor), även om de inte är lika högnivåinriktade och koncisa som Triton eller cuTile, var rimliga substitut. Dessutom erbjuder CuTeDSL obegränsad åtkomst till maskinvaruprimitiver, vilket ger oss den kontrollnivå vi behöver för att uppnå topprestanda i alla specialfall. Med CUDA var möjligheterna till kompileringsspecialisering begränsade, eftersom varje extra parameter tillförde en oskälig kostnad i form av kompileringstider eller komplexitet för byggsystemet. Med ett Just-In-Time-kompilerat DSL kan vi däremot specialisera oss mer aggressivt utan att offra iterationstakten. Till skillnad från TVM IR, som använde CUDA som ett mellanliggande steg, fann vi att den MLIR-baserade pipelinen som skapade PTX var mer robust och enklare att felsöka än en som genererade CUDA som ett mellanliggande steg. Genom att använda anpassade operationer för att kapsla in PTX-instruktioner upptäckte vi dessutom att vi till fullo kunde utnyttja alla funktioner i den underliggande maskinvaran.

Huvuddelen av kompileringskostnaden för CUDA-kärnor läggs vanligtvis på mallexpansionsfasen för kärnor, eftersom optimering och kompilering från PTX-stadiet och framåt kan slutföras på några sekunder. Med många kompileringsparametrar kan detta kompileringssteg explodera till flera minuter för en icke-trivial kärna. Denna flaskhals accentueras av det faktum att möjligheter till parallellkompilerings uppstår först senare, på IR- och PTX-nivåerna, medan C++-mallhantering sker sekventiellt. Det finns visserligen knep för att komma till rätta med detta genom att flytta mallinstansieringar till separata kompileringsenheter, men de kräver betydande finjusteringar av byggsystem och användning av otympliga makron. Däremot kräver CuTeDSL-kärnor som kompileras med Just-In-Time (JIT) inte komplexa byggregler och hämmar inte kompileringstiderna, eftersom utvecklingscykeln kan fokuseras på en enda punkt i konfigurationsutrymmet, medan noggranna enhetstester senare kan täcka in kärnans hela designutrymme.

Även om JIT-kompilering flyttar kostnader från utvecklingstid till körtid är omkostnaden mer än acceptabel. När en modell driftsätts aktiverar den bara ett fåtal punkter från hela konfigurationsutrymmet, vilket medför omkostnader i storleksordningen sekunder under uppstart/uppvärmning av inferensmotorn. Däremot kan förkompilering av alla konfigurationer för en kärna ta flera minuter. Dessutom kan omkostnaderna mildras ytterligare genom att de JIT-kompilerade artefakterna cachas.

CUDAs tunga kompileringskostnad förvärras av svårigheten att spåra ursprunget till fel när mallbaserad dispatch och specialisering används, vilket är fallet med CuTe och CUTLASS. Med ett DSL kan fel pekas ut på den kodrad som genererade det, vilket underlättar felsökningen avsevärt. Dessutom är det mycket enklare att mappa den genererade koden till det DSL-uttryck som genererade den via mellanliggande MLIR-steg eller genom att förlita sig på felsökningsinformation som är kopplad till den genererade PTX:en. Tack vare dessa aspekter av DSL:et upptäckte vi att vi var mycket mer produktiva när vi utvecklade kärnor.

Även om kärnor har många likheter mellan modeller måste de ändå utrustas med en rad konfigurationsreglage som kan påverka prestandan. Relevanta parametrar inkluderar:

Dold dimension / Huvuddimension / Antal attention-huvuden

Indata-/utdatadatatyper

Kvantiseringsschema

Kompileringskonstanter (bias, aktiveringsfunktioner med mera)

ROSE-arkitektur, inklusive dess motor, modell, kärnbibliotek och tjänster.

Kompileringsspecialisering är användbart för både prestanda och felsökning. Särskilt i en pipelined processor, vilket är fallet med GPU:er, är en icke-existerande gren bättre än en gren som alltid eller aldrig tas. Dessutom leder kompilering av en kärna för en specifik uppsättning konfigurationsreglage till färre PTX-/SASS-instruktioner, som är enklare att analysera och felsöka när man använder cuda-gdb under en undersökning.

För en given batchstorlek och konfiguration kan den förväntade trafiken till DRAM och antalet aritmetiska operationer uppskattas. Därefter kan optimala beslut fattas från fall till fall för optimal gridallokering, synkronisering och användning av vissa maskinvarufunktioner. CuTeDSL tillåter oss att definiera och återanvända komponenter för djup specialisering utan att det medför en oacceptabel kompileringsomkostnad.

I allmänhet måste prestandan hos inferenskärnor optimeras för två separata användningsfall:

Prefill: liten batchstorlek (få sjok, sekvenser), många tokens per sekvens. Genomströmningen ska maximeras för både inbäddningar och indatatokens för språkmodeller. Latens på tiotals till hundratals millisekunder.

Avkodning (Decode): små till stora batchstorlekar, en (ingen spekulativ avkodning) till ett fåtal tokens per sekvens (draft-target-spekulation, prediktion med flera tokens). Latens nära kärnstartens latens.

För GEMM och attention körs prefill-stadiet på ett stort antal tokens och prestandan begränsas av enhetens beräkningsgenomströmning. Avkodning innebär samma operation på små batchstorlekar, där exekveringstiderna domineras av kostnaden för att läsa in aktiveringar och vikter från DRAM till de strömmande multiprocessorerna och tensor-kärnorna, vilket istället begränsas av enhetens minnesbandbredd. Även om batchstorleken eller det totala antalet tokens i en batch är en relativt tydlig indikator är övergångströskeln mellan prefill och avkodning mer subtil eftersom minnesbandbredden också bestäms av modellens dolda dimension och aktiveringens datatyp. Kärnor delas därför upp i prefill- och avkodningsimplementeringar om algoritmiska skillnader eller grid-/blocktilldelningar skulle gynna någon av dem.

Den parameter som genomgående skiljer sig åt mellan prefill- och avkodningskärnor är antalet warps i ett block (antalet trådar dividerat med warp-storleken på 32 trådar per warp). Det är önskvärt att göra denna parameter till en kompileringskonstant eftersom det kan underlätta registerallokering. I CUDA görs detta vanligtvis genom att ange antalet trådar via __launch_bounds__ och tillhandahålla antalet warps som en mallparameter:

cpp

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

}

Antalet register som är tillgängliga för alla trådar i ett block är en maskinvarubegränsning, nämligen 64K på Hopper och Blackwell. Om antalet trådar är känt i förväg vet kompilatorn exakt hur många register den kan tilldela varje tråd och kan lyckas kompilera mer aggressivt optimerade kärnor. Detta har dock ett pris i form av kompileringstid, eftersom kärnan nu måste ompileras för olika warp-storlekar. Med CuTeDSL är detta inte längre ett problem, eftersom vi inte behöver betala kompileringskostnaden för hela konfigurationsutrymmet i förväg. DSL:et förväntar sig också att blockstorleken är ett konstant uttryck, vilket implicit utför den optimering som antyds av __latch_bounds__ i CUDA.

Förutom att göra antalet warps konstant är det önskvärt att variera det mellan prefill och avkodning. För att till fullo utnyttja minnesbandbredden i en GPU räcker det inte med att alltid tilldela ett högt antal warps per block, eftersom full bandbredd endast kan uppnås genom att använda så många som möjligt av de tillgängliga strömmande multiprocessorerna (SM). Följaktligen väljer vi, för ett givet antal tokens och kända datatyper och dolda dimensioner, först en idealisk vektordimension per tråd (8 för bfloat16, 4 för float32) för att säkerställa att den genererade koden använder 128-bitars läs-/skrivinstruktioner (256-bitars på Blackwell). När vi känner till vektorstorleken kan vi beräkna det antal warps som krävs av kärnan. Vi försöker först maximera antalet SM: om det finns färre warps än SM väljs en enda warp per block. Detta är vanligtvis avkodningsvarianten av kärnan. Om det finns tillräckligt med warps ökar vi gradvis antalet warps till 8, 16 eller 32 och benchmarkar en övre gräns efter vilken vi vanligtvis inte längre ser någon prestandaförbättring.

Medan Triton utförde några av dessa optimeringar under huven upptäckte vi att vi kan uppnå liknande resultat genom att noggrant använda CuTe-layoutalgebra, utan att begränsas av bristerna hos en högnivåkompilator. Särskilt när horisontella reduktioner under warp-nivå är inblandade fann vi att CuTeDSL erbjöd den finkorniga kontroll som krävs för att generera optimal kod.

I vissa situationer, för att distribuera arbete över så många block som möjligt, kan cooperative_groups::this_grid().sync() användas för att synkronisera alla deltagande block. Vissa begränsningar gäller: eftersom ett block inte kan avplaneras från en SM när det väl har börjat exekveras, och för att garantera framsteg framåt, får antalet block inte överstiga antalet SM som är tillgängliga på enheten. Detta säkerställer att alla block när som helst kan uppdatera och avläsa de globala minnesplatser som är involverade i synkroniseringen. Denna begränsning upprätthålls av cudaLaunchCooperativeKernel. Kostnaden för en grid-barriär är ungefär 3 mikrosekunder på Hopper, plus kostnaden för att rensa alla tidigare minnesoperationer via en implicit minnesbarriär i vissa situationer. Sammantaget är detta jämförbart med omkostnaderna för att dela upp en kärna i två starter, vilket leder till ett prestandakänsligt designbeslut. Trots det faktum att Programmatic Dependent Launch (PDL) kan minska startomkostnaderna från och med Hopper är det värt att undvika omkostnaderna för avkodningskärnor.

Om en av arbetsbelastningarna kräver fler block än antalet SM bör kärnan delas upp, varvid man förlitar sig på den imlicita barriären vid kärnstarter. Detta gäller i första hand för prefill. Om det antal block som är idealiskt för alla arbetsbelastningar får plats i griden föredras en grid-barriär, vilket är fallet med avkodningskärnor som arbetar med ett mindre antal tokens.

I CuTeDSL kan dessa beslut fattas i startkoden på värdsidan, och kärnorna kan sammansättas och varvas med grid-barriärer efter behov, vilket illustreras i exemplet nedan. Med en sådan hybrid strategi kan avkodningskärnor optimeras för att uppnå minimal latens och prefill-kärnor kan optimeras för maximal genomströmning. Till skillnad från CUDA-baserade lösningar kan denna logik definieras mer koncist och med lägre kompileringsomkostnader.

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

När det gäller våra MoE all-to-all-kärnor, vilket beskrivs nedan, hjälpte flexibiliteten att bygga kärnor med eller utan grid-barriärer och optimera dem vid kompileringstid oss att uppnå avsevärda prestandavinster genom CuTeDSL.

Den första kärnan vi byggde i CuTeDSL var en specialiserad RMS-normkärna som optimerats för små dimensioner. QK-normen, som används i attention-lagret i de senaste Qwen-modellerna, tillämpas per huvud på 64, 128 eller 256 element. Till skillnad från vanlig RMS-norm, som arbetar på dolda dimensioner som överstiger tusental och kräver en horisontell reduktion genom både delat minne och register, behöver QK-normen endast reduktion över warps, med möjlighet att klämma in flera huvud i en enda warp. Detta blev ett användningsfall där vi kunde optimera aggressivt vid kompileringstid för att hitta de bästa tilldelningarna från token till warp för alla huvuddimensioner och datatyper.

RMS-norm är en ganska okomplicerad operation, även om den kräver en horisontell summa för att beräkna kvadratroten ur normen för en vektor:

cpp

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

Eftersom attention-huvuden vanligtvis är små och trådoperationer är starkt vektoriserade kan en enda warp beräkna kvadratnormen. För att maximera minnesbandbredden bör varje tråd bearbeta 128 bitars indata, vilket motsvarar 8 bfloat16-värden. Följaktligen räcker en halv warp för att bearbeta ett huvud, och kvadratnormen kan bestämmas med hjälp av en horisontell reduktion som utförs av shuffle_sync_bfly. Med utgångspunkt i huvuddimensionen och datatypen kan vi fastställa antalet huvud per warp och specificera argumenten samt antalet reduktionsrundor vid kompileringstillfället.

Vi kan sedan bygga en specialiserad kärna med relativt få rader kod, genom att förlita oss på automatiskt vektoriserade in-/utläsningar och använda arkitekturspecifika operationer på lägre nivå för warp-reduktionen. Kärnan avgör vilket huvud (Q eller K) den ska arbeta med och identifierar motsvarande vikter som ska tillämpas. Den laddar en hel vektor från huvudet och summerar först lokalt. Därefter beräknas kvadratnormen baserat på antalet trådar per huvud. Slutligen bestäms RMS-normen, automatiskt vektorerad, och skrivs tillbaka till minnet. Modellspecifika parametrar som WEIGHT_BIAS och EPS är konstantvikta.

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)

Jämfört med en vanlig RMS-normkärna ger denna specialiserade variant 2–3 gånger bättre genomströmning och latens för både prefill och avkodning.

MoE dispatch och combine blandar om indatatokens till de expertmodeller de skickas till, och utför en viktad ackumulering efter att MoE GEMM-kärnorna har utfört de relevanta matrismultiplikationerna. Prestandan är av särskilt intresse eftersom olika produkter förlitar sig på MoE-modeller i olika storlekar. Beroende på driftsättningstyp skiljer vi på olika kärnor som ska optimeras:

TP=1: Lokal kärna för en enskild GPU. Dispatch blandar om tokens till relevanta expertmodeller och combine ackumulerar dem. Används främst för Qwen3 30B-driftsättningar på H100.

DP=1, TP ≠ 1: Tensor-parallel, enodskärna. Indatan replikeras, dispatch väljer ut tokens för den aktuella rangordningen och combine ackumulerar dem över alla rangordningar. Ackumulering över antingen NVLink eller NVLS. Används för stora prefill-körningar.

DP ≠ 1, NVLink: Dataparallell kärna med valfri tensorparallellism. Varje rangordning skickar tokens till alla andra rangordningar som hostar specifika expertmodeller. Routinginformation aggregeras, tokens överförs via NVIDIA NVLink. Används för dataparallella H200- eller GB200-driftsättningar.

DP ≠ 1, InfiniBand: Peer-to-Peer-dataparallell och tensorparallell kärna som körs över både NVLink och NVIDIA Quantum InfiniBand. Används för flernods-Hopper-driftsättningar.

Den lokala kärnan räknar antalet tokens som tilldelats varje expertmodell baserat på routingtabellen och tilldelar ett unikt index till varje token i listan över tokens som hör till respektive expertmodell. Efter routingsteget kopieras tokens till buffertar som kan fyllas ut (padding) för att skickas vidare till en MoE GEMM-kärna. Med hjälp av tokenplatserna utför combine-kärnan sedan en ackumulering i register och skriver de slutliga tokens till en utdatabuffert.

Dispatch-fasen varierar mellan prefill och avkodning: för få tokens finns det lite routinginformation att läsa ut, så vi bearbetar den i ett enda block och signalerar resten av griden, där enskilda delblock kopierar sjok av tokens. Tokentilldelning bestäms genom att atomärt inkrementera räknare i det delade minnet. Eftersom prefill måste bearbeta mer routinginformation delas routingen upp över flera block i en separat kärna. Index tilldelas först lokalt med hjälp av räknare i det delade minnet och aggregeras sedan över alla block i det globala minnet för att bestämma positionen för varje token. Griden för denna kärna begränsas till antalet SM för att kunna använda grid-omfattande barriärer. En efterföljande kärna utför kopieringarna, startade över ett mycket större antal block och warps för att till fullo utnyttja den tillgängliga minnesbandbredden.

Combine-kärnan är relativt enkel eftersom den utför vektorerad ackumulering, med stöd av routinginformationen och de förskjutningar som beräknats av dispatch. CuTeDSL visade sig vara värdefullt här, eftersom beräkningen kunde uttryckas genom enkla CuTe-tensorer som kompileras till effektiv kod för alla datatyper och dolda dimensioner.

I en tensorparallell driftsättning replikeras routinginformationen och indatatokens. Den lokala dispatch-kärnan kan återanvändas fullt ut med en liten ändring för att routa och endast välja ut de tokens som tilldelats de expertmodeller som finns på den aktuella rangordningen. Med vissa parametrar och kompileringsspecialisering återanvänds de lokala dispatch-kärnorna.

Eftersom antalet rangordningar vanligtvis är lägre än antalet expertmodeller som en token routas till kan man i genomsnitt anta att varje token har minst en av sina repliker routad till varje rangordning. I denna situation kan den lokala combine-kärnan återanvändas för att aggregera de lokala kopiorna av en token, följt av en all-reduce för att summera över de deltagande rangordningarna. När det gäller prefill-kärnan startar vi separata kärnor, men avkodning slår däremot samman båda operationerna till en enda start.

För all-reduce använder vi vanlig NVLink om endast två rangordningar deltar, och byter till NVLS när 4 eller fler rangordningar är inblandade för en ökad genomströmning på 10–20 %. Tack vare DSL:et kan vi dela den mesta av reduceringslogiken och peer-buffertutbytena mellan MoE och de vanliga all-reduce-kärnorna. Detta är en situation där nära kontroll över maskinvaran är fördelaktig, eftersom DSL-operationer över multimem.ld_reduce och multimem.st krävs för NVLS. Implementeringen av själva reduktionen följer i stor utsträckning exemplet all_reduce från NCCL-testerna.

De dataparallella kärnorna förlitar sig på NVLink för att aggregera tokens för flera enheter. I dispatch-fasen, innan tokens utbyts, routar noder först lokalt för att identifiera målinformationen för förskjutning på varje rangordning för varje token de skickar. Därefter börjar de kopiera över tokens till mottagningsbuffertar som har allokerats för varje avsändare för att undvika konflikter. Skrivningar med release-store per token används för att signalera slutförande. På mottagarsidan aggregeras routinginformationen för att kunna kondensera alla tokens till en enda buffert för MoE GEMM. Mottagare avläser den flagga som skrivits av peern med en acquire-load för att vänta på att skrivningarna ska slutföras.

På combine-sidan kopieras tokens över NVLink till källrangordningarna, med en flagga per token som indikerar mottagande. Combine-kärnan liknar de lokala och enodbaserade varianterna, med små skillnader i adressering och viss extra synkroniseringslogik.

Vi byggde ett bibliotek med PTX-omslutningar (wrappers) för att exponera release-acquire-läsningar och -skrivningar begränsade till både system- och enhetsnivå för att implementera denna kärna. Genom att förbättra granulariteten hos synkroniseringsflaggor och allokera fler block för att skicka och ta emot tokens över NVLink uppnådde vi dubbelt så bra resultat som med våra tidigare CUDA-kärnor.

Våra Peer-to-Peer-kärnor uppnådde tidigare banbrytande prestanda på Hopper-GPU:er i kombination med ConnectX-7 InfiniBand-adaptrar. Vi har sedan dess portat combine-kärnorna till CuTeDSL för ytterligare prestandavinster. Eftersom CUDA-versionerna redan i stor utsträckning förlitade sig på kompileringsspecialisering visade de sig vara kompileringsflaskhalsar i vårt kodarkiv, vilket tvingade oss att utesluta vissa parametrar. Genom att migrera till det JIT-kompilerade DSL:et specialiserade vi oss för fler parametrar och finjusterade gridallokeringar, vilket introducerade ytterligare skillnader mellan prefill och avkodning för en latensförbättring på omkring 3 %.

Ur ROSE-kärnornas perspektiv delas Blackwell-stödet in i två kategorier: arkitekturoberoende och arkitekturspecifika kärnor. Arkitekturspecifika kärnor, som förlitar sig på tensor-kärnor, kräver explicit arbete för att anpassas. Tack vare användningen av DSL anpassas de arkitekturoberoende kärnorna däremot utan ansträngning.

En Blackwell-GPU har upp till 160 SM, till skillnad från Hoppers 132. För de flesta kärnor är detta en kompileringskonstant, och logiken i startkoden på värdsidan tar hänsyn till detta för att avgöra vilken optimal kärna (prefill kontra avkodning) som ska köras. Dessutom introducerar Blackwell 256-bitars läsinstruktioner, som enkelt kan genereras genom att öka vektorvidden (den sista dimensionen i CuTe-tensorer) för att förbättra utnyttjandet av minnesbandbredden.

Blackwell introducerar hårdvarustöd för ett betydande antal kvantiseringsscheman (MXFP8, MXFP4, NVFP4). Stöd för dessa implementerades främst genom att finjustera befintliga kärnor för FP8-blockvis skalning. JIT-stöd är avgörande här, eftersom nya lägen kan stödjas utan att konfigurationsutrymmet för kompilering exploderar. Tidigare fann vi att Triton var svårhanterligt när det gäller att uttrycka horisontella reduktioner över ett fåtal element, över en halv eller en fjärdedel av en warp. Finkornigare kontroll över reduktioner var avgörande för att uppnå optimal prestanda vid beräkning av ett absolut maximalt värde över respektive 16 eller 32 element för NVFP4 och MXFP8/MXFP4.

Eftersom det är ett relativt nytt språk hade språkmodeller (LLM) haft lite exponering för CuTeDSL-kod under sin förträning. Dessutom skilde sig stilen på vår nuvarande kod, med kärnor optimerade för aktiveringar och normer, något från de befintliga CUTLASS- eller Quack-exemplen, som var kraftigt optimerade för TMA- och arkitekturspecifika matrismultiplikationsinstruktioner. Att nå en punkt där språkmodeller kunde hjälpa till vid kärnskrivning var en gradvis process.

Vi byggde upp en ursprunglig korpus av kärnor och verktygsfunktioner manuellt, enligt konsekventa mönster, innan vi kunde förlita oss på agenter för kärnkonvertering. Allt eftersom korpusen växte kunde vi utföra grovöversättningar från CUDA till CuTeDSL, med betydande manuellt efterarbete som följd. I synnerhet upptäckte vi att modellerna kämpade med konverteringar från pekar till tensor och med vektorisering, och fastnade i försök att åtgärda typfel när målet var att få ett test att kompileras. De har visat liten förståelse för kärnornas gridbaserade natur, och förväxlat tråd- och gridsynkronisering. Mest frustrerande var frestelsen att hallucinera icke-existerande översättningar av CUDA-specifika bibliotek i stället för att bygga mot våra egna CuTeDSL-ersättningar. Med tiden, i takt med att korpusen av CuTeDSL-kärnor växte, upptäckte vi att språkmodeller kan bygga det mesta av skelettkoden och fylla i luckorna, vilket gör att vi kan fokusera på övergripande design.

Jämfört med CUDA upptäcker vi att CuTeDSL levererar mycket bättre kompileringsprestanda, vilket gör att vi kan specialisera oss för fler parametrar och uppnå prestandaförbättringar. Även om kärnorna är något mer ordrika på grund av avsaknaden av vissa abstraktioner jämfört med Triton, anser vi att komplexiteten är motiverad eftersom finkornigare kontroll över maskinvaran gör att vi kan bygga snabbare kärnor.

Vid skrivandets stund är prestandagapet mellan CUTLASS C++- och CuTeDSL-kärnor minimalt, vilket öppnar upp spännande möjligheter att komponera och slå samman högt presterande GEMM-kärnor med kommunikation mellan enheter, normalisering eller aktiveringar. Genom att satsa på detta DSL kommer vi att kunna utvinna ytterligare prestanda ur våra enheter för att förbättra genomströmningen över våra arbetsbelastningar, från inbäddningar till språkmodeller.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Inferens på Perplexity AI

Varför CuTeDSL?

Utvecklarupplevelse

Stort konfigurationsutrymme

Förifyllning och avkodningsspecialisering

Gridsynkronisering

CuTeDSL-kärnor

QK-norm

Mixture-of-Experts Dispatch/Combine

Lokal

Enod

Dataparallell

Peer-to-Peer

Blackwell-stöd

AI-assisterad kärngenerering

Sammanfattning och framtida arbete

Referenser

Citering