CuTeDSL presso Perplexity

In Perplexity utilizziamo il Runtime-Optimized Serving Engine (ROSE) interno per gestire modelli che spaziano dagli embedding agli LLM con trilioni di parametri. Posto dietro tutte le API, inclusi Sonar, Search ed Embeddings, il nostro runtime è stato creato per consentire ai ricercatori di esplorare i modelli e le architetture più adatti, distribuendo al contempo rapidamente i migliori candidati in produzione. Per ottenere ciò, dobbiamo essere in grado di implementare, distribuire e ottimizzare l'inferenza per architetture all'avanguardia a un ritmo rapido. In questo post del blog, esploriamo come i kernel CuTeDSL integrati nel nostro motore di inferenza ci consentano di creare i kernel GPU necessari per portare i modelli alle massime prestazioni sulle GPU NVIDIA.

Inferenza presso Perplexity AI

I modelli personalizzati, creati dai vari team di Perplexity, sono ospitati internamente su GPU NVIDIA Hopper e Blackwell utilizzando il motore di inferenza ROSE creato dal team AI Inference. ROSE è stato inizialmente creato per servire meglio i modelli Llama personalizzati sia per la decodifica del linguaggio che per la classificazione, con un'interfaccia compatibile con NVIDIA Triton Inference Server. Nel tempo, si è evoluto in un motore di inferenza completo che serve LLM sostanziali, oltre a modelli di ranking, classificazione, scoring e altri modelli basati su transformer per una vasta gamma di team per alimentare i prodotti di Perplexity. L'obiettivo di progettazione di ROSE è essere semplice: nel tempo, viene investito impegno per raggiungere le massime prestazioni, ma la priorità principale è adattarsi alle architetture e ai casi d'uso in continua evoluzione per i nostri modelli.

Espandi

A livello generale, ROSE consiste in un motore che adatta un modello a un'interfaccia con cui i client possono interagire. I motori sono altamente riutilizzabili, gestiscono la pianificazione e il raggruppamento delle richieste (batching) e recuperano token decodificati, embedding o punteggi dai modelli sottostanti. L'inizializzazione dei dispositivi, i canali di comunicazione inter-processo e inter-nodo e il caricamento dei pesi sono tutti gestiti da essi. Per gli LLM, il motore è responsabile del batching, del chunking e del campionamento. È anche responsabile dell'allocazione dello storage KV, sia per l'attention completa che lineare, nonché del prefix matching. I motori di embedding sono più semplici, poiché eseguono solo il batching online. L'esecuzione è ottimizzata in modo tale che il motore prepari il round di lavoro successivo mentre il modello è in esecuzione sull'acceleratore in parallelo.

ROSE espone una ricca collezione di layer personalizzati per la modellazione, avvolgendo kernel GPU creati con CuTeDSL, Triton, CUDA, CUTLASS e cuBLAS. Per le implementazioni di moltiplicazione di matrici e attention, ci affidiamo principalmente ai kernel NVIDIA da CUTLASS e cuBLAS. Il resto della pipeline, inclusi i kernel di embedding, norm, MoE e activation, fa parte della libreria di kernel di ROSE. Per supportare al meglio l'hardware su cui eseguiamo l'inferenza, ROSE è dotato di proprie primitive di comunicazione inter-dispositivo per il routing, il dispatch e il combine MoE. Sebbene gli altri kernel siano relativamente semplici, devono essere approfonditamente specializzati per ottenere le massime prestazioni su tutte le diverse dimensioni nascoste e sui dispositivi su cui vengono eseguiti. Inoltre, facciamo affidamento sulla specializzazione in fase di compilazione per adattarci a tutte le variabili e le modifiche di vari modelli (weight bias rispetto a no weight bias per RMS norm, funzioni di attivazione, ecc.). Di conseguenza, abbiamo deciso di utilizzare i DSL per implementarli: inizialmente abbiamo creato ROSE utilizzando Triton. Tuttavia, da allora abbiamo adottato CuTeDSL come ambiente principale per la programmazione GPU.

Perché CuTeDSL?

Da una prospettiva matematica, i kernel di inferenza sono solitamente poco sofisticati. Di conseguenza, durante lo sviluppo di un kernel, la complessità generalmente non risiede nell'esprimere quale calcolo debba essere eseguito, ma nel come esattamente l'hardware debba farlo. Per creare i migliori kernel, richiediamo un linguaggio che offra le astrazioni corrette rispetto all'hardware sottostante. Il linguaggio giusto ci consente di specificare le operazioni di alto livello per eseguire il calcolo, specializzandole e ottimizzandole in modo aggressivo in tutti i punti in un ampio Spazio di configurazione.

CuTeDSL è un DSL basato su Python per GPU NVIDIA costruito sopra l'algebra dei layout CuTe e MLIR. Viene compilato just-in-time in PTX altamente ottimizzato, pur mantenendo il controllo sulle primitive hardware di basso livello. Mira a offrire le stesse funzionalità e prestazioni di CUTLASS C++, senza i costosi tempi di compilazione che ostacolano la velocità di sviluppo. Tra i linguaggi che abbiamo esaminato o con cui abbiamo lavorato (inclusi CUDA, cuTile, TVM IR e Triton), lo consideriamo la soluzione migliore per lo sviluppo dei nostri kernel.

Complessivamente, sebbene non così di alto livello e concisi come Triton o cuTile, abbiamo riscontrato che l'algebra dei layout CuTe e le sue composizioni (esposte tramite cute.Tensor) fossero sostituti ragionevoli. Inoltre, CuTeDSL offre accesso illimitato alle primitive hardware, fornendoci il livello di controllo richiesto per ottenere le massime prestazioni in tutti i casi limite. Con CUDA, le opportunità di specializzazione in fase di compilazione erano limitate, poiché ogni parametro aggiuntivo aggiungeva un costo irragionevole ai tempi di compilazione o complessità al sistema di build. Al contrario, con un DSL compilato Just-In-Time, possiamo specializzare in modo più aggressivo senza sacrificare il ritmo di iterazione. A differenza di TVM IR, che utilizzava CUDA come passaggio intermedio, abbiamo trovato la pipeline basata su MLIR che riduceva a PTX più robusta e più facile da eseguire il debug rispetto a quella che emetteva CUDA come passaggio intermedio. Inoltre, con l'uso di operazioni personalizzate per avvolgere le istruzioni PTX, abbiamo scoperto di poter sfruttare appieno tutte le funzionalità dell'hardware sottostante.

Esperienza dello sviluppatore

La maggior parte del costo di compilazione dei kernel CUDA viene solitamente spesa nella fase di espansione dei template dei kernel, poiché l'ottimizzazione e la compilazione dallo stadio PTX in poi possono essere completate in pochi secondi. Con molti parametri in fase di compilazione, questa fase di compilazione può esplodere in diversi minuti per un kernel non banale. Questo collo di bottiglia è accentuato dal fatto che le opportunità per la compilazione parallela si presentano solo più tardi, ai livelli IR e PTX, con la gestione dei template C++ che avviene in modo sequenziale. Sebbene esistano trucchi per risolvere questo problema, spostando le istanze di template in unità di compilazione separate, richiedono modifiche significative ai sistemi di build e l'uso di macro ingombranti. Al contrario, i kernel CuTeDSL compilati Just-In-Time (JIT) non richiedono complesse regole di build e non ostacolano i tempi di compilazione, poiché il ciclo di sviluppo può essere focalizzato su un singolo punto nello Spazio di configurazione, mentre accurati test unitari in seguito possono coprire l'intero spazio di progettazione del kernel.

Anche se la compilazione JIT sposta i costi dal tempo di sviluppo al runtime, il sovraccarico è più che accettabile. Quando un modello viene distribuito, attiva solo alcuni dei punti dall'intero Spazio di configurazione, comportando sovraccarichi su scala di secondi durante l'avvio/riscaldamento del motore di inferenza. Al contrario, la pre-compilazione di tutte le configurazioni di un kernel potrebbe richiedere diversi minuti. Inoltre, i sovraccarichi possono essere ulteriormente mitigati memorizzando nella cache gli artifact compilati JIT.

L'elevato costo di compilazione di CUDA è accentuato dalla difficoltà di tracciare l'origine degli errori quando vengono utilizzati il dispatch e la specializzazione basati su template, come nel caso di CuTe e CUTLASS. Con un DSL, gli errori possono essere indicati nella riga di codice che li ha generati, semplificando significativamente il debug. Inoltre, è molto più semplice mappare il codice generato all'istruzione DSL che lo ha originato tramite stadi MLIR intermedi o facendo affidamento sulle informazioni di debug allegate al PTX generato. Grazie a questi aspetti del DSL, ci siamo trovati molto più produttivi durante lo sviluppo dei kernel.

Ampio Spazio di configurazione

Sebbene i kernel condividano molte similitudini tra i modelli, devono comunque essere dotati di una serie di manopole di configurazione che possono avere un impatto sulle prestazioni. I parametri rilevanti includono:

Dimensione nascosta/Dimensione della testa/Numero di teste di attention

Tipi di dati di input/output

Schema di quantizzazione

Costanti in fase di compilazione (bias, funzioni di attivazione, ecc.)

Espandi

La specializzazione in fase di compilazione è utile sia per le prestazioni che per il debug. In particolare in un processore pipelined, come nel caso delle GPU, un ramo inesistente è migliore di un ramo sempre intrapreso o mai intrapreso. Inoltre, compilare un kernel per un set specifico di manopole di configurazione si traduce in un minor numero di istruzioni PTX/SASS, che sono più facili da analizzare ed eseguire il debug quando si utilizza cuda-gdb durante un'indagine.

Per una determinata dimensione del batch e configurazione, è possibile stimare il traffico previsto verso la DRAM e il numero di operazioni aritmetiche. Successivamente, è possibile prendere decisioni ottimali caso per caso per l'allocazione ottimale della griglia, la sincronizzazione e l'uso di determinate funzionalità hardware. CuTeDSL ci consente di definire e riutilizzare componenti per una specializzazione profonda senza incorrere in un sovraccarico inaccettabile in fase di compilazione.

Specializzazione di precompilazione e decodifica

Generalmente, le prestazioni dei kernel di inferenza devono essere ottimizzate per due casi d'uso separati:

Prefill: dimensione del batch ridotta (pochi chunk, sequenze), molti token per sequenza. throughput da massimizzare sia per gli embedding che per i token di input degli LLM. Latenza da decine a centinaia di millisecondi.

Decode: dimensioni del batch da piccole a grandi, da uno (nessuna decodifica speculativa) a pochi token per sequenza (speculazione target-draft, previsione multi-token). Latenza vicina alla latenza di avvio del kernel.

Per GEMM e attention, la fase di prefill viene eseguita su un gran numero di token e le prestazioni sono vincolate dal throughput di calcolo del dispositivo. La decodifica comporta la stessa operazione su piccole dimensioni del batch, con tempi di esecuzione dominati dal costo di caricamento delle attivazioni e dei pesi dalla DRAM ai multiprocessori di streaming e ai core tensor, vincolati invece dalla larghezza di banda della memoria del dispositivo. Anche se la dimensione del batch o il numero totale di token in un batch è un indicatore relativamente chiaro, la soglia di passaggio tra prefill e decode è più sottile perché la larghezza di banda della memoria è determinata anche dalla dimensione nascosta del modello e dal tipo di dati dell'attivazione. I kernel vengono conseguentemente suddivisi in implementazioni prefill e decode se differenze algoritmiche o assegnazioni di griglia/blocco avvantaggerebbero uno di essi.

Il parametro che è universalmente diverso tra i kernel prefill e decode è il numero di warp in un blocco (il numero di thread, diviso per la dimensione del warp di 32 thread per warp). È auspicabile rendere questo parametro una costante in fase di compilazione, poiché può aiutare nell'allocazione dei registri. In CUDA, ciò viene solitamente fatto specificando il numero di thread tramite __launch_bounds__ e fornendo il numero di warp come parametro template:

Il numero di registri disponibili per tutti i thread in un blocco è un vincolo hardware, 64K su Hopper e Blackwell. Se il numero di thread è noto in anticipo, il compilatore sa esattamente quanti registri può assegnare a ogni thread e può riuscire a compilare kernel ottimizzati in modo più aggressivo. Tuttavia, questo comporta un costo in termini di tempo di compilazione, poiché ora il kernel deve essere ricompilato per diverse dimensioni di warp. Con CuTeDSL, questo non è più un problema, poiché non dobbiamo pagare il costo della compilazione per l'intero Spazio di configurazione in anticipo. Il DSL prevede anche che la dimensione del blocco sia un'espressione costante, eseguendo implicitamente l'ottimizzazione implicata da __latch_bounds__ in CUDA.

Oltre a rendere costante il numero di warp, è auspicabile variarlo tra prefill e decode. Per utilizzare appieno la larghezza di banda della memoria di una GPU, non è sufficiente assegnare sempre un numero elevato di warp per blocco, poiché la piena larghezza di banda può essere raggiunta solo utilizzando il maggior numero possibile di Streaming Multiprocessors (SM) disponibili. Di conseguenza, per un dato numero di token e tipi di dati e dimensioni nascoste noti, scegliamo innanzitutto una dimensione vettoriale ideale per thread (8 per bfloat16, 4 per float32) per garantire che il codice generato utilizzi le istruzioni di caricamento/memorizzazione a 128 bit (256 bit su Blackwell). Conoscendo la dimensione del vettore, possiamo calcolare il numero di warp richiesti dal kernel. Proviamo innanzitutto a massimizzare il numero di SM: se ci sono meno warp rispetto agli SM, viene scelto un singolo warp per blocco. Questa è solitamente la variante decode del kernel. Se ci sono abbastanza warp, aumentiamo gradualmente il numero di warp a 8, 16 o 32, effettuando il benchmarking di un limite superiore dopo il quale solitamente non vediamo più un miglioramento nelle prestazioni.

Sebbene Triton abbia eseguito alcune di queste ottimizzazioni sotto il cofano, abbiamo riscontrato che possiamo ottenere risultati simili utilizzando attentamente l'algebra dei layout CuTe, senza essere vincolati dai difetti di un compilatore di alto livello. In particolare, quando sono coinvolte riduzioni orizzontali sub-warp, abbiamo riscontrato che CuTeDSL ci offriva il controllo granulare richiesto per generare codice ottimale.

Sincronizzazione della griglia

In alcune situazioni, per distribuire il lavoro tra il maggior numero possibile di blocchi, è possibile utilizzare cooperative_groups::this_grid().sync() per sincronizzare tutti i blocchi partecipanti. Si applicano delle restrizioni: poiché un blocco non può essere de-pianificato da un SM una volta iniziata l'esecuzione, per garantire il progresso in avanti, il numero di blocchi non può superare il numero di SM disponibili sul dispositivo, per garantire che in qualsiasi momento tutti i blocchi possano aggiornare e sottoporre a polling le posizioni di memoria globale coinvolte nella sincronizzazione. Questa restrizione è applicata da cudaLaunchCooperativeKernel. Il costo di una barriera di griglia è di circa 3us su Hopper, più il costo di svuotamento di tutte le precedenti operazioni di memoria tramite una barriera di memoria implicita in alcune situazioni. Complessivamente, questo è paragonabile al sovraccarico di dividere un kernel in due lanci, portando a una decisione di progettazione sensibile alle prestazioni. Nonostante il fatto che il Programmatic Dependent Launch (PDL) possa ridurre i sovraccarichi di lancio da Hopper in poi, per i kernel decode, evitare il sovraccarico è utile.

Se uno dei carichi di lavoro richiede più blocchi rispetto al numero di SM, il kernel dovrebbe essere suddiviso, facendo affidamento sulla barriera implicita dei lanci del kernel. Questo è principalmente il caso del prefill. Se il numero di blocchi ideale per tutti i carichi di lavoro rientra nella griglia, viene preferita una barriera di griglia, come nel caso dei kernel decode che operano su un numero inferiore di token.

In CuTeDSL, queste decisioni possono essere prese nel codice di lancio lato host e i kernel possono essere composti e intercalati con barriere di griglia secondo necessità, come illustrato nell'esempio seguente. Con una tale strategia ibrida, i kernel decode possono essere ottimizzati per raggiungere latenze minime e quelli prefill possono essere ottimizzati per il massimo throughput. A differenza delle soluzioni basate su CUDA, questa logica può essere definita in modo più conciso e con meno sovraccarico in fase di compilazione.

Nel caso dei nostri kernel MoE all-to-all, come dettagliato di seguito, avere la flessibilità di creare kernel con o senza barriere di griglia e ottimizzarli in fase di compilazione ci ha aiutato a ottenere notevoli guadagni di prestazioni tramite CuTeDSL.

Kernel CuTeDSL

QK Norm

Il primo kernel che abbiamo creato in CuTeDSL è stato un kernel RMS norm specializzato, ottimizzato per dimensioni ridotte. QK norm, utilizzato nel layer di self-attention dei recenti modelli Qwen, viene applicato su base per-testa, su 64, 128 o 256 elementi. A differenza del normale RMS norm, che opera su dimensioni nascoste che superano le migliaia, richiedendo una riduzione orizzontale attraverso sia la memoria condivisa che i registri, QK norm necessita di riduzione solo tra i warp, con la possibilità di comprimere più teste in un singolo warp. Questo è diventato un caso d'uso in cui potevamo ottimizzare in modo aggressivo in fase di compilazione per trovare le migliori assegnazioni token-to-warp per tutte le dimensioni di testa e tipi di dati.

RMS norm è un'operazione abbastanza semplice, sebbene richieda una somma orizzontale per calcolare la norma al quadrato di un vettore:

Poiché le teste di attention sono solitamente piccole e le operazioni per-thread sono pesantemente vettorializzate, un singolo warp può calcolare la norma quadrata. Per massimizzare la larghezza di banda della memoria, ogni thread dovrebbe elaborare 128 bit di input, il che si traduce in 8 valori bfloat16. Di conseguenza, un mezzo-warp è sufficiente per elaborare una testa e la norma al quadrato può essere determinata utilizzando una riduzione orizzontale eseguita da shuffle_sync_bfly. Partendo dalla dimensione della testa e dal tipo di dati, possiamo determinare il numero di teste per warp e specificare gli argomenti e il numero di round di riduzione al momento della compilazione.

Possiamo quindi creare un kernel specializzato con relativamente poche righe di codice, facendo affidamento su carichi/memorizzazioni vettorializzati automaticamente e utilizzando operazioni specifiche dell'architettura di livello inferiore per la riduzione dei warp. Il kernel determina su quale testa (Q o K) lavorare e identifica i pesi corrispondenti da applicare. Carica un intero vettore dalla testa e somma innanzitutto localmente. Successivamente, in base al numero di thread per testa, viene calcolata la norma al quadrato. Infine, l'RMS norm, vettorializzato automaticamente, viene determinato e riscritto in memoria. I parametri specifici del modello, come WEIGHT_BIAS e EPS, sono constant-folded.

Rispetto a un normale kernel RMS norm, questa variante specializzata offre prestazioni 2-3 volte migliori in termini di throughput e latenza sia per il prefill che per il decode.

Mixture-of-Experts Dispatch/Combine

MoE dispatch e combine mischiano i token di input agli esperti a cui sono instradati, eseguendo un'accumulazione ponderata dopo che i kernel MoE GEMM hanno eseguito le relative moltiplicazioni di matrici. Le prestazioni sono di particolare interesse, poiché vari prodotti si affidano a modelli MoE di varie dimensioni. In base ai tipi di distribuzione, distinguiamo diversi kernel da ottimizzare:

TP=1: Kernel locale per una singola GPU. Dispatch mischia i token agli esperti rilevanti e combine li accumula. Utilizzato principalmente per distribuzioni Qwen3 30B su H100.

DP=1, TP ≠ 1: Kernel parallelo ai tensor, a nodo singolo. L'input viene replicato, il dispatch seleziona i token per il rank corrente e il combine li accumula tra tutti i rank. Accumulazione tramite NVLink o NVLS. Utilizzato per prefiller di grandi dimensioni.

DP ≠ 1, NVLink: Kernel parallelo ai dati con parallelismo tensor opzionale. Ogni rank invia token a ogni altro rank che ospita esperti specifici. Le informazioni di instradamento vengono aggregate, i token vengono trasferiti tramite NVIDIA NVLink. Utilizzato per distribuzioni H200 o GB200 parallele ai dati.

DP ≠ 1, InfiniBand: Kernel parallelo ai dati Peer-to-Peer e parallelo ai tensor che opera sia su NVLink che su NVIDIA Quantum InfiniBand. Utilizzato per distribuzioni Hopper multi-nodo.

Locale

Il kernel locale conta il numero di token assegnati a ogni esperto in base alla tabella di instradamento e assegna un indice univoco a ogni token all'interno dell'elenco di token appartenenti a ogni esperto. Dopo il passaggio di instradamento, i token vengono copiati in buffer facoltativamente riempiti (padded) per essere passati a un kernel MoE GEMM. Utilizzando le posizioni dei token, il kernel combine esegue quindi un'accumulazione nei registri, scrivendo i token finali in un buffer di output.

La fase di dispatch varia tra prefill e decode: per pochi token, ci sono poche informazioni di instradamento da leggere, quindi le elaboriamo in un singolo blocco e segnaliamo il resto della griglia, con singoli blocchi che copiano chunk di token. L'assegnazione dei token è determinata dall'incremento atomico dei contatori nella memoria condivisa. Poiché il prefill deve elaborare più informazioni di instradamento, l'instradamento viene suddiviso su più blocchi in un kernel separato. Gli indici vengono prima assegnati localmente utilizzando i contatori nella memoria condivisa e vengono poi aggregati tra tutti i blocchi nella memoria globale per determinare la posizione di ogni token. La griglia di questo kernel è bloccata sul numero di SM per poter utilizzare barriere a livello di griglia. Un kernel successivo esegue le copie, lanciato su un numero molto maggiore di blocchi e warp per utilizzare appieno la larghezza di banda della memoria disponibile.

Il kernel combine è relativamente semplice, poiché esegue un'accumulazione vettorializzata, facendo affidamento sulle informazioni di instradamento e sugli offset calcolati dal dispatch. CuTeDSL si è rivelato prezioso qui, poiché il calcolo poteva essere espresso attraverso semplici tensori CuTe che si compilano in codice efficiente per tutti i tipi di dati e le dimensioni nascoste.

Nodo singolo

In una distribuzione parallela ai tensor, le informazioni di instradamento e i token di input vengono replicati. Il kernel di dispatch locale può essere interamente riutilizzato con una leggera modifica per instradare e selezionare solo i token assegnati agli esperti situati sul rank corrente. Con alcuni parametri e la specializzazione in fase di compilazione, i kernel di dispatch locale vengono riutilizzati.

Poiché il numero di rank è solitamente inferiore al numero di esperti a cui viene instradato un token, nel caso medio si può presumere che ogni token abbia almeno una delle sue repliche instradata su ogni rank. In questa situazione, il kernel combine locale può essere riutilizzato per aggregare le copie locali di un token, seguite da un all-reduce per sommare tra i rank partecipanti. Nel caso del kernel di prefill, lanciamo kernel separati, tuttavia il decode fonde entrambe le operazioni in un unico lancio.

Per all-reduce, utilizziamo il semplice NVLink se partecipano solo due rank, passando a NVLS quando sono coinvolti 4 o più rank per un aumento del 10-20% del throughput. Grazie al DSL, possiamo condividere la maggior parte della logica di riduzione e gli scambi di buffer peer tra MoE e i normali kernel all-reduce. Questa è una situazione in cui un controllo ravvicinato sull'hardware è vantaggioso, poiché le operazioni DSL su multimem.ld_reduce e multimem.st sono richieste per NVLS. L'implementazione della riduzione stessa segue in gran parte l'esempio all_reduce dai test NCCL.

Data Parallel

I kernel Data Parallel si affidano a NVLink per aggregare i token per più dispositivi. Nella fase di dispatch, prima di scambiare i token, i peer instradano innanzitutto localmente per identificare l'offset target su ogni rank per ogni token che stanno inviando. Successivamente, iniziano a copiare i token nei buffer di ricezione allocati per ogni mittente per evitare conflitti. Le scritture release-store per-token vengono utilizzate per segnalare il completamento. Sul lato ricevente, le informazioni di instradamento vengono aggregate per poter condensare tutti i token in un unico buffer per MoE GEMM. I ricevitori interrogano il flag scritto dal peer con un acquire-load per attendere il completamento delle scritture.

Sul lato combine, i token vengono copiati tramite NVLink sui rank di origine, con un flag per-token che indica la ricezione. Il kernel combine è simile alle varianti locale e a nodo singolo, con lievi differenze nell'indirizzamento e un'ulteriore logica di sincronizzazione.

Abbiamo creato una libreria di wrapper PTX per esporre i caricamenti e le memorizzazioni release-acquire con ambito sia a livello di sistema che di dispositivo per implementare questo kernel. Migliorando la granularità dei flag di sincronizzazione e allocando più blocchi per inviare e ricevere token tramite NVLink, abbiamo ottenuto miglioramenti 2x rispetto ai nostri precedenti kernel CUDA.

Peer-to-Peer

I nostri kernel Peer-to-Peer avevano precedentemente raggiunto prestazioni all'avanguardia su GPU Hopper abbinate ad adattatori InfiniBand ConnectX-7. Da allora abbiamo trasferito i kernel combine su CuTeDSL per ulteriori guadagni in termini di prestazioni. Poiché le versioni CUDA facevano già molto affidamento sulla specializzazione in fase di compilazione, si sono rivelate i colli di bottiglia in termini di tempo di compilazione nel nostro repository, costringendoci a omettere alcuni parametri. Migrando al DSL compilato JIT, abbiamo specializzato per più parametri e modificato le allocazioni della griglia, introducendo ulteriori differenze tra prefill e decode per circa un aumento del 3% della latenza.

Supporto Blackwell

Dal punto di vista dei kernel ROSE, il supporto Blackwell rientra in due categorie: kernel agnostici rispetto all'architettura e kernel specifici per l'architettura. I kernel specifici per l'architettura, basati sui core Tensor, richiedono un lavoro esplicito per l'adattamento. Tuttavia, grazie all'uso del DSL, i kernel agnostici rispetto all'architettura si adattano in modo banale.

Una GPU Blackwell ha fino a 160 SM, in contrasto con i 132 di Hopper. Per la maggior parte dei kernel, questa è una costante in fase di compilazione e la logica nel codice del launcher lato host ne tiene conto per determinare il kernel (prefill vs decode) ottimale da eseguire. Inoltre, Blackwell introduce anche istruzioni di caricamento a 256 bit, che possono essere facilmente emesse aumentando la larghezza del vettore (ultima dimensione dei tensori CuTe) per migliorare l'utilizzo della larghezza di banda della memoria.

Blackwell introduce il supporto hardware per un numero significativo di schemi di quantizzazione (MXFP8, MXFP4, NVFP4). Il supporto per questi è stato implementato principalmente modificando i kernel esistenti per il ridimensionamento a blocchi FP8. Il supporto JIT è cruciale qui, poiché nuove modalità possono essere supportate senza far esplodere lo Spazio di configurazione in fase di compilazione. In passato, abbiamo trovato Triton scomodo nell'esprimere riduzioni orizzontali su pochi elementi, oltre metà o un quarto di un warp. Un controllo più granulare sulle riduzioni è stato cruciale per raggiungere prestazioni ottimali durante il calcolo di un valore massimo assoluto su 16 o 32 elementi rispettivamente per NVFP4 e MXFP8/MXFP4, all'interno dello Spazio.

Generazione di kernel assistita dall'IA

Essendo un linguaggio relativamente nuovo, gli LLM hanno avuto poca esposizione al codice CuTeDSL durante il loro pre-addestramento. Inoltre, lo stile del nostro codice attuale, con kernel ottimizzati per attivazioni e norme, era leggermente diverso dagli esempi CUTLASS o Quack esistenti, che erano pesantemente ottimizzati per TMA e istruzioni di moltiplicazione di matrici specifiche dell'architettura. Raggiungere un punto in cui gli LLM potessero assistere nella scrittura dei kernel è stato un processo graduale.

Abbiamo creato manualmente un corpus iniziale di kernel e funzioni di utilità, seguendo schemi coerenti, prima di poter fare affidamento sugli agenti per la conversione del kernel. Man mano che il corpus progrediva, potevamo eseguire traduzioni approssimative da CUDA a CuTeDSL, con una significativa rielaborazione manuale a seguire. In particolare, abbiamo riscontrato che i modelli faticavano con le conversioni da puntatore a tensore e la vettorializzazione, bloccandosi nel tentativo di correggere gli errori di tipo quando miravano a far compilare un test. Hanno mostrato poca comprensione della natura basata sulla griglia dei kernel, confondendo la sincronizzazione tra thread e griglia. Cosa più frustrante, c'era la tentazione di allucinare traduzioni inesistenti di librerie specifiche per CUDA invece di costruire rispetto alle nostre sostituzioni CuTeDSL. Nel tempo, man mano che il corpus di kernel CuTeDSL cresceva, abbiamo riscontrato che gli LLM possono costruire la maggior parte del codice scheletro e riempire gli spazi vuoti, permettendoci di concentrarci sulla progettazione di alto livello.

Conclusione e lavoro futuro

Rispetto a CUDA, riscontriamo che CuTeDSL offre prestazioni in fase di compilazione molto migliori, consentendoci di specializzare per più parametri e ottenere miglioramenti delle prestazioni. Sebbene i kernel siano leggermente più prolissi a causa della mancanza di alcune astrazioni rispetto a Triton, riscontriamo che la complessità è giustificata poiché un controllo più granulare sull'hardware ci consente di creare kernel più veloci.

Al momento della scrittura, il divario di prestazioni tra i kernel CUTLASS C++ e CuTeDSL è minimo, aprendo interessanti opportunità per comporre e fondere kernel GEMM ad alte prestazioni con comunicazione inter-dispositivo, normalizzazione o attivazioni. Impegnandoci in questo DSL, saremo in grado di estrarre ulteriori prestazioni dai nostri dispositivi per migliorare il throughput su tutti i nostri carichi di lavoro, dagli embedding agli LLM.

Riferimenti

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/

Citazione

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