CuTeDSL a Perplexity

A Perplexity, fem servir el motor de servei optimitzat per a temps d'execució (ROSE) desenvolupat internament per gestionar models que van des de representacions vectorials (embeddings) fins a LLMs d'un bilió de paràmetres. Situat darrere de totes les APIs, incloent-hi Sonar, Search i Embeddings, el nostre entorn d'execució es va crear per permetre als investigadors explorar els models i arquitectures més adequats, alhora que desplega ràpidament els millors candidats a producció. Per aconseguir-ho, hem de ser capaços de implementar, desplegar i optimitzar la inferència per a arquitectures d'última generació a un ritme ràpid. En aquesta entrada del bloc, explorem com els nuclis de CuTeDSL integrats al nostre motor d'inferència ens permeten construir els nuclis de GPU necessaris per portar els models al seu rendiment màxim en GPUs de NVIDIA.

Els models personalitzats, creats per diversos equips de Perplexity, s'allotgen internament en GPUs NVIDIA Hopper i Blackwell utilitzant el motor d'inferència ROSE creat per l'equip d'inferència d'IA. ROSE es va construir inicialment per oferir un millor servei als models Llama personalitzats tant per a la descodificació de llenguatge com per a la classificació, amb una interfície compatible amb NVIDIA Triton Inference Server. Amb el temps, es va convertir en un motor d'inferència complet que serveix LLMs importants, així com models de ranking, classificació, puntuació i altres models basats en transformadors per a una àmplia gamma d'equips per potenciar els productes de Perplexity. L'objectiu de disseny de ROSE és ser senzill: amb el temps, s'inverteix esforç per assolir el màxim rendiment, però la prioritat principal és adaptar-se a les arquitectures i casos d'ús en constant evolució dels nostres models.

L'espai de configuració del nucli entre tipus de dades, dimensions del cap, constants de temps de compilació i paràmetres d'execució.

A un nivell general, ROSE consta d'un motor que adapta un model a una interfície amb la qual els clients poden interactuar. Els motors són altament reutilitzables, gestionen la planificació i el processament per lots de les sol·licituds, i recuperen els tokens descodificats, les representacions vectorials o les puntuacions dels models subjacents. La inicialització de dispositius, els canals de comunicació entre processos i entre nodes, i la càrrega de pesos són gestionats per ells. Per als LLMs, el motor s'encarrega de la creació de lots, la fragmentació i el mostreig. També s'encarrega de l'assignació d'emmagatzematge KV, tant per a l'atenció completa com lineal, així com de la coincidència de prefixos. Els motors de representació vectorial són més senzills, ja que només realitzen lots en línia. L'execució està optimitzada de manera que el motor prepara la següent ronda de treball mentre el model s'executa a l'accelerador en paral·lel.

ROSE ofereix una àmplia col·lecció de capes personalitzades per a la modelització, encapsulant nuclis de GPU construïts amb CuTeDSL, Triton, CUDA, CUTLASS i cuBLAS. Per a les implementacions de multiplicació de matrius i atenció, confiem principalment en els nuclis de NVIDIA de CUTLASS i cuBLAS. La resta de la canonada, incloent-hi els nuclis d'incrustació, normalització, MoE i activació, formen part de la biblioteca de nuclis de ROSE. Per donar el millor suport al maquinari on executem la inferència, ROSE està equipat amb les seves pròpies primitives de comunicació entre dispositius per a l'encaminament, despatx i combinació de MoE. Tot i que els altres nuclis són relativament senzills, s'han d'especialitzar a fons per aconseguir el rendiment màxim en totes les diferents dimensions ocultes i dispositius on s'executen. A més, confiem en l'especialització en temps de compilació per adaptar-nos a tots els ajustos i configuracions de diversos models (biaix de pes enfront de cap biaix de pes per a la norma RMS, funcions d'activació, etc.). En conseqüència, vam decidir utilitzar DSL per implementar-los: inicialment vam construir ROSE utilitzant Triton. No obstant això, des de llavors hem adoptat CuTeDSL com a entorn principal per a la programació de GPU.

Des d'una perspectiva matemàtica, els nuclis d'inferència solen ser poc sofisticats. En conseqüència, en desenvolupar un nucli, la complexitat generalment no rau a expressar quina computació s'ha de dur a terme, sinó en com exactament el maquinari l'ha de fer. Per construir els millors nuclis, necessitem un llenguatge que ofereixi les abstraccions adequades sobre el maquinari subjacent. El llenguatge adequat ens permet especificar les operacions d'alt nivell per realitzar la computació, especialitzant-les i optimitzant-les de manera agressiva en tots els punts d'un ampli espai de configuració.

CuTeDSL és un DSL basat en Python per a GPUs NVIDIA construït sobre l'àlgebra de disposicions (layout algebra) de CuTe i MLIR. Es compila just a temps (just-in-time) a un PTX altament optimitzat, alhora que manté el control sobre les primitives de maquinari de baix nivell. Pretén oferir les mateixes característiques i rendiment que CUTLASS C++, sense els costosos temps de compilació que frenen la velocitat de desenvolupament. Dels llenguatges que hem analitzat o amb els quals hem treballat (incloent-hi CUDA, cuTile, TVM IR i Triton), considerem que és el més adequat per al desenvolupament dels nostres nuclis.

En general, tot i no ser tan d'alt nivell i concís com Triton o cuTile, vam trobar que l'àlgebra de disposicions de CuTe i les seves composicions (exposades a través de cute.Tensor) eren substituts raonables. A més, CuTeDSL ofereix accés sense restriccions a les primitives de maquinari, proporcionant-nos el nivell de control que necessitem per assolir el rendiment màxim en tots els casos límit. Amb CUDA, les oportunitats d'especialització en temps de compilació eren limitades, ja que cada paràmetre addicional afegia un cost inacceptable als temps de compilació o complexitat al sistema de construcció. En canvi, amb un DSL compilat just a temps, podem especialitzar-nos de manera més agressiva sense sacrificar el ritme d'iteració. A diferencia de TVM IR, que utilitzava CUDA com a pas intermedi, vam trobar que la canonada basada en MLIR que baixava a PTX era més robusta i fàcil de depurar que una que emetia CUDA com a pas intermedi. A més, amb l'ús d'operacions personalitzades per encapsular les instruccions PTX, vam descobrir que podíem aprofitar al màxim totes les característiques del maquinari subjacent.

La major part del cost de compilació dels nuclis CUDA es dedica normalment a la fase d'expansió de plantilles dels nuclis, ja que l'optimització i la compilació des de l'etapa PTX en endavant es poden completar en qüestió de segons. Amb molts paràmetres en temps de compilació, aquesta etapa de compilació pot disparar-se a diversos minuts per a un nucli no trivial. Aquest coll d'ampolla s'accentua pel fet que les oportunitats de compilació paral·lela només sorgeixen més tard, als nivells IR i PTX, amb la gestió de plantilles C++ succeint de manera seqüencial. Tot i que existeixen trucs per solucionar-ho, movent les instanciacions de plantilles a unitats de compilació separades, requereixen modificacions significatives als sistemes de construcció i l'ús de macros feixugues. En contrast, els nuclis CuTeDSL que es compilen Just-In-Time (JIT) no requereixen regles de construcció complexes ni obstaculitzen els temps de compilació, ja que el cicle de desenvolupament es pot centrar en un sol punt de l'espai de configuració, mentre que unes proves unitàries acurades posteriorment poden cobrir tot el disseny espai del nucli.

Tot i que la compilació JIT trasllada els costos del temps de desenvolupament al temps d'execució, la sobrecàrrega és més que acceptable. Quan es desplega un model, aquest activa només uns pocs punts de tot l'espai de configuració, generant sobrecàrregues de l'ordre de segons durant l'inici/escalfament del motor d'inferència. En canvi, precompilar totes les configuracions d'un nucli podria abastar diversos minuts. A més, les sobrecàrregues es poden mitigar encara més emmagatzemant en memòria cau els artefactes compilats per JIT.

El pesat cost de compilació de CUDA s'accentua per la dificultat de rastrejar l'origen dels errors quan s'utilitza l'enviament basat en plantilles i l'especialització, com és el cas de CuTe i CUTLASS. Amb un DSL, els errors es poden apuntar a la línia de codi que els ha generat, facilitant significativament la depuració. A més, és molt més senzill assignar el codi generat a la sentència DSL que l'ha originat a través d'etapes intermediàries de MLIR o basant-se en la informació de depuració adjunta al PTX generat. Gràcies a aquests aspectes del DSL, ens hem trobat molt més productius a l'hora de desenvolupar nuclis.

Tot i que els nuclis comparteixen moltes similituds entre models, encara s'han d'equipar amb un rang de botons de configuració que poden tenir un impacte en el rendiment. Els paràmetres rellevants inclouen:

Dimensió oculta / Dimensió de capçalera / Nombre de capçals d'atenció

Tipus de dades d'entrada/sortida

Esquema de quantificació

Constants en temps de compilació (biaixos, funcions d'activació, etc.)

Arquitectura ROSE, incloent-hi el seu motor, model, biblioteca del nucli i serveis.

L'especialització en temps de compilació és útil tant per al rendiment com per a la depuració. Particularment en un processador en canonada (pipelined), com és el cas de les GPUs, una branca inexistent és millor que una branca sempre presa o mai presa. A més, compilar un nucli per a un conjunt específic de paràmetres de configuració dóna lloc a menys instruccions PTX/SASS, que són més fàcils d'analitzar i depurar quan s'utilitza cuda-gdb durant una investigació.

Per a una mida de lot i una configuració determinades, es pot estimar el trànsit esperat a DRAM i el nombre d'operacions aritmètiques. Posteriorment, es poden prendre decisions òptimes cas per cas per a una assignació òptima de la graella, la sincronització i l'ús de determinades característiques del maquinari. CuTeDSL ens permet definir i reutilitzar components per a una especialització profunda sense incórrer en una sobrecàrrega inacceptable en el temps de compilació.

En general, el rendiment dels nuclis d'inferència s'ha d'optimitzar per a dos casos d'ús diferents:

Preompliment (Prefill): mida de lot petita (pocs fragments, seqüències), molts tokens per seqüència. El rendiment s'ha de maximitzar tant per a les incrustacions com per als tokens d'entrada dels LLMs. Latència de desenes a centenars de mil·lisegons.

Descodificació (Decode): mides de lot de petites a grans, d'un (sense descodificació especulativa) a uns quants tokens per seqüència (especulació d'objectiu esborrany, predicció de múltiples tokens). Latència propera a la latència de llançament del nucli.

Per a GEMM i l'atenció, l'etapa de preompliment s'executa en un gran nombre de tokens i el rendiment està limitat pel rendiment de còmput del dispositiu. La descodificació implica la mateixa operació en mides de lot petites, amb temps d'execució dominats pel cost de carregar les activacions i els pesos de la DRAM als multiprocessadors de flux i nuclis tensorials, limitats en canvi per l'amplada de banda de memòria del dispositiu. Tot i que la mida del lot o el nombre total de tokens en un lot és un indicador relativament clar, el llindar de canvi entre el preompliment i la descodificació és més subtil perquè l'amplada de banda de la memòria també ve determinada per la dimensió oculta del model i el tipus de dada de l'activació. En conseqüència, els nuclis es divideixen en implementacions de preompliment i descodificació si les diferències algorítmiques o les assignacions de graella/bloc beneficiarien a qualsevol d'ells.

El paràmetre que és universalment diferent entre els nuclis de preompliment i descodificació és el nombre de grups de fils o warps en un bloc (el nombre de fils, dividit per la mida del warp de 32 fils per warp). És desitjable fer d'aquest paràmetre una constant en temps de compilació, ja que pot ajudar en l'assignació de registres. A CUDA, això es fa normalment especificant el nombre de fils mitjançant __launch_bounds__ i proporcionant el nombre de warps com a paràmetre de plantilla:

cpp

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

}

El nombre de registres disponibles per a tots els fils d'un bloc és una restricció de maquinari, 64K a Hopper i Blackwell. Si es coneix el nombre de fils per avançat, el compilador sap exactament quants registres pot assignar a cada fil i pot aconseguir compilar nuclis optimitzats de manera més agressiva. No obstant això, això comporta un cost en temps de compilació, ja que ara el nucli s'ha de recompilar per a diferents mides de warp. Amb CuTeDSL, això ja no és un problema, ja que no hem de pagar el cost de compilació de tot l'espai de configuració per endavant. El DSL també espera que la mida del bloc sigui una expressió constant, realitzant implícitament l'optimització implicada per __latch_bounds__ a CUDA.

A més de fer que el nombre de warps sigui una constant, és desitjable variar-lo entre el preompliment i la descodificació. Per aprofitar al màxim l'amplada de banda de memòria d'una GPU, no n'hi ha prou amb assignar sempre un nombre elevat de warps per bloc, ja que l'amplada de banda completa només es pot aconseguir utilitzant tants Multiprocessadors de Flux (SM) disponibles com sigui possible. En conseqüència, per a un nombre determinat de tokens i tipus de dades i dimensions ocultes coneguts, primer triem una dimensió de vector per fil ideal (8 per a bfloat16, 4 per a float32) per assegurar-nos que el codi generat utilitzi les instruccions de càrrega/emmagatzematge de 128 bits (256 bits a Blackwell). Coneixent la mida del vector, podem calcular el nombre de warps requerits pel nucli. Primer intentem maximitzar el nombre de SM: si hi ha menys warps que SM, s'escull un sol warp per bloc. Aquesta és normalment la variant de descodificació del nucli. Si hi ha prou warps, augmentem gradualment el nombre de warps a 8, 16 o 32, avaluant un límit superior a partir del qual normalment ja no veiem una millora en el rendiment.

Tot i que Triton va realitzar algunes d'aquestes optimitzacions sota el capó, vam descobrir que podem aconseguir resultats similars utilitzant amb cura l'àlgebra de disposicions de CuTe, sense estar limitats per les mancances d'un compilador d'alt nivell. Particularment quan es tracta de reduccions horitzontals sub-warp, vam trobar que CuTeDSL ens ofereix el control de graella fina necessari per generar codi òptim.

En algunes situacions, per distribuir el treball entre tants blocs com sigui possible, es pot utilitzar cooperative_groups::this_grid().sync() per sincronitzar tots els blocs participants. S'apliquen restriccions: com que un bloc no es pot desprogramar d'un SM un cop comença l'execució, per garantir el progrés cap endavant, el nombre de blocs no pot superar el nombre de SM disponibles al dispositiu, per assegurar que en qualsevol moment tots els blocs puguin actualitzar i consultar les ubicacions de memòria global implicades en la sincronització. Aquesta restricció la reforça cudaLaunchCooperativeKernel. El cost d'una barrera de graella és d'aproximadament 3us a Hopper, més el cost de buidar totes les operacions de memòria anteriors mitjançant una barrera de memòria implícita en algunes situacions. En general, això és comparable a la sobrecàrrega de dividir un nucli en dos llançaments, la qual cosa porta a una decisió de disseny sensible al rendiment. Malgrat el fet que el llançament dependent programàtic (PDL) pot reduir les sobrecàrregues de llançament des de Hopper en endavant, per als nuclis de descodificació, val la pena evitar la sobrecàrrega.

Si un de les càrregues de treball requereix més blocs que el nombre de SM, el nucli s'ha de dividir, basant-se en la barrera implícita dels llançaments de nuclis. Aquest és principalment el cas del preompliment. Si el nombre de blocs que són ideals per a totes les càrregues de treball s'ajusta a la graella, es prefereix una barrera de graella, com és el cas dels nuclis de descodificació que operen amb un nombre més petit de tokens.

A CuTeDSL, aquestes decisions es poden prendre en el codi de llançament del costat de l'amfitrió i els nuclis es poden compondre i entrellaçar amb barreres de graella segons sigui necessari, tal com il·lustra l'exemple següent. Amb aquesta estratègia híbrida, els nuclis de descodificació es poden optimitzar per assolir latències mínimes i els de preompliment es poden optimitzar per obtenir el màxim rendiment. A diferència de les solucions basades en CUDA, aquesta lògica es pot definir de manera més concisa i amb menys sobrecàrrega en el temps de compilació.

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

En el cas dels nostres nuclis MoE all-to-all, tal com es detalla a continuació, tenir la flexibilitat de construir nuclis amb o sense barreres de graella i optimitzar-los en el moment de la compilació ens va ajudar a aconseguir guanys substancials de rendiment mitjançant CuTeDSL.

El primer nucli que vam construir a CuTeDSL va ser un nucli de norma RMS specialized, optimitzat per a dimensions petites. La norma QK, utilitzada a la capa d'atenció pròpia (self-attention) dels models Qwen recents, s'aplica per capçalera, en 64, 128 o 256 elements. A diferència de la norma RMS regular, que opera en dimensions ocultes que superen els milers, requerint una reducció horitzontal tant a través de la memòria compartida com dels registres, la norma QK només necessita reducció a través dels warps, amb la possibilitat d'ajustar múltiples capçals en un sol warp. Això es va convertir en un cas d'ús on podíem optimitzar agressivament en el moment de la compilació per trobar les millors assignacions de token a warp per a totes les dimensions de capçal i tipus de dades.

La norma RMS és una operació bastant directa, tot i que requereix una suma horitzontal per calcular la norma alquadrada d'un vector:

cpp

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

Com que els capçals d'atenció solen ser petits i les operacions per fil estan fortament vectoritzades, un sol warp pot calcular la norma alquadrada. Per maximitzar l'amplada de banda de la memòria, cada fil ha de processar 128 bits d'entrada, la qual cosa es tradueix en 8 valors bfloat16. En conseqüència, un mig warp (half-warp) és suficient per processar un capçal, i la norma alquadrada es pot determinar mitjançant una reducció horitzontal realitzada per shuffle_sync_bfly. Partint de la dimensió del capçal i del tipus de dada, podem determinar el nombre de capçals per warp i especificar els arguments i el nombre de rondes de reducció en el moment de la compilació.

A continuació, podem construir un nucli especialitzat amb relativament poques línies de codi, basant-nos en càrregues/emmagatzematges vectoritzats automàticament i utilitzant operacions específiques d'arquitectura de nivell inferior per a la reducció del warp. El nucli determina en quin capçal (Q o K) ha de treballar i identifica els pesos corresponents a aplicar. Carrega un vector sencer del capçal i primer suma localment. Posteriorment, segons el nombre de fils per capçal, es calcula la norma alquadrada. Finalment, es determina la norma RMS, vectoritzada automàticament, i s'escriu de nou a la memòria. Els paràmetres específics del model, com ara WEIGHT_BIAS i EPS, es pleguen a constants (constant-folded).

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)

En comparació amb un nucli de norma RMS regular, aquesta variant especialitzada ofereix un rendiment i una latència entre 2 i 3 vegades superiors tant per al preompliment com per a la descodificació.

El despatx i la combinació de MoE reordenen els tokens d'entrada als experts als quals s'encaminen, realitzant una acumulació ponderada després que els nuclis GEMM de MoE realitzin les multiplicacions de matrius pertinents. El rendiment és d'un interès particular, ja que diversos productes depenen de models MoE de diverses mides. Segons els tipus de desplegament, distingim diferents nuclis a optimitzar:

TP=1: Nucli local per a una sola GPU. El despatx reordena els tokens als experts pertinents i la combinació els acumula. S'utilitza principalment per a desplegaments de Qwen3 30B a H100.

DP=1, TP ≠ 1: Nucli de node únic amb paral·lelisme de tensors. L'entrada es replica, el despatx selecciona els tokens per al rang actual i la combinació els acumula a tots els rangs. Acumulació a través de NVLink o NVLS. S'utilitza per a preomplidors grans.

DP ≠ 1, NVLink: Nucli paral·lel de dades amb paral·lelisme de tensors opcional. Cada rang envia tokens a la resta de rangs que allotgen experts específics. La informació d'encaminament s'agrega, i els tokens es transfereixen a través d'NVIDIA NVLink. S'utilitza per a desplegaments H200 o GB200 paral·lels de dades.

DP ≠ 1, InfiniBand: Nucli paral·lel de dades i paral·lelisme de tensors peer-to-peer que opera tant a través d'NVLink com d'NVIDIA Quantum InfiniBand. S'utilitza per a desplegaments Hopper multinode.

El nucli local compta el nombre de tokens assignats a cada expert basant-se en la taula d'encaminament i assigna un índex únic a cada token dins de la llista de tokens que pertanyen a cada expert. Després del pas d'encaminament, els tokens es copien en memòries intermèdies (buffers) amb farciment opcional (padded) per passar-los a un nucli GEMM de MoE. Utilitzant les ubicacions dels tokens, el nucli de combinació realitza aleshores una acumulació en registres, escrivint els tokens finals en una memòria intermèdia de sortida.

La fase de despatx varia entre el preompliment i la descodificació: per a pocs tokens, hi ha poca informació d'encaminament per llegir, així que la processem en un sol bloc i enviem un senyal a la resta de la graella, amb blocs individuals copiant fragments de tokens. L'assignació de tokens es determina incrementant de manera atòmica els comptadors de la memòria compartida. Com que el preompliment ha de processar més informació d'encaminament, l'encaminament es divideix en diversos blocs en un nucli separat. Els índexs s'assignen primer localment utilitzant comptadors de la memòria compartida i després s'agreguen a tots els blocs de la memòria global per determinar la posició de cada token. La graella d'aquest nucli es limita (clamped) al nombre de SM per poder utilitzar barreres a nivell de graella. Un nucli posterior realitza les còpies, llançat a través d'un nombre molt més gran de blocs i warps per aprofitar al màxim l'amplada de banda de memòria disponible.

El nucli de combinació és relativament senzill, ja que realitza una acumulació vectoritzada, basant-se en la informació d'encaminament i en els desplaçaments (offsets) calculats pel despatx. CuTeDSL va demostrar ser valuós aquí, ja que el còmput es podia expressar a través de simples tensors CuTe que es compilen en codi eficient per a tots els tipus de dades i dimensions ocultes.

En un desplegament paral·lel de tensors, la informació d'encaminament i els tokens d'entrada es repliquen. El nucli de despatx local es pot reutilitzar completament amb un lleuger canvi per encaminar i seleccionar només els tokens assignats als experts ubicats al rang actual. Amb alguns paràmetres i especialització en temps de compilació, es reutilitzen els nuclis de despatx locals.

Com que el nombre de rangs sol ser inferior al nombre d'experts als quals s'encamina un token, en el cas mitjà es pot assumir que cada token té almenys una de les seves rèpliques encaminada a cada rang. en aquesta situació, el nucli de combinació local es pot reutilitzar per agregar les còpies locals d'un token, seguit d'un all-reduce per sumar entre els rangs participants. En el cas del nucli de preompliment, llancem nuclis separats, no obstant això, la descodificació fusiona ambdues operacions en un sol llançament.

Per a all-reduce, utilitzem NVLink pla si només hi participen dos rangs, passant a NVLS quan s'hi inclouen 4 o més rangs per a un augment del 10 al 20% en el rendiment. Gràcies al DSL, podem compartir gran part de la lògica de reducció i els intercanvis de memòria intermèdia entre iguals (peer buffer exchanges) entre els nuclis MoE i els de all-reduce normals. Aquesta és una situació on un control estret sobre el maquinari és beneficiós, ja que es requereixen operacions DSL sobre multimem.ld_reduce i multimem.st per a NVLS. La implementació de la reducció en si segueix en gran mesura l'exemple all_reduce de les proves de NCCL.

Els nuclis paral·lels de dades es basen en NVLink per agregar tokens per a múltiples dispositius. En la fase de despatx, abans d'intercanviar tokens, els iguals (peers) primer s'encaminen localment per identificar el desplaçament objectiu de cada rang per a cada token que estan enviant. Posteriorment, comencen a copiar els tokens a les memòries intermèdies de recepció assignades per a cada remitent per evitar conflictes. S'utilitzen escriptures de tipus release-store per token per indicar la finalització. Al costat del receptor, la informació d'encaminament s'agrega per poder condensar tots els tokens en una sola memòria intermèdia per a GEMM de MoE. Els receptors consulten el senyal escrit pel peer amb un acquire-load per esperar que les escriptures es completin.

Pel costat de la combinació, els tokens es copien a través d'NVLink als rangs d'origen, amb un senyal per token que indica la recepció. El nucli de combinació és similar a les variants locals i de node únic, amb lleugeres diferències en l'adreçament i alguna lògica de sincronització addicional.

Vam construir una biblioteca de microprogramari (wrappers) de PTX per exposar càrregues i emmagatzematges release-acquire amb àmbit tant a nivell de sistema com de dispositiu per implementar aquest nucli. Millorant la granularitat dels senyals de sincronització i assignant més blocs per enviar i rebre tokens a través d'NVLink, vam aconseguir millores de 2x respecte als nostres nuclis CUDA anteriors.

Els nostres nuclis Peer-to-Peer havien aconseguit prèviament un rendiment d'última generació en GPUs Hopper combinades amb adaptadors InfiniBand ConnectX-7. Des de llavors, hem portat els nuclis de combinació a CuTeDSL per obtenir més guanys de rendiment. Com que les versions de CUDA ja depenien molt de l'especialització en temps de compilació, van demostrar ser els colls d'ampolla en el temps de compilació al nostre repositori, obligant-nos a ometre alguns paràmetres. En migrar al DSL compilat per JIT, ens vam especialitzar per a més paràmetres i vam ajustar les assignacions de graella, introduint més diferències entre el preompliment i la descodificació per a un impuls d'aproximadament un 3% en la latència.

Des de la perspectiva dels nuclis de ROSE, el suport per a Blackwell es divideix en dues categories: nuclis agnòstics d'arquitectura i nuclis específics d'arquitectura. Els nuclis específics d'arquitectura, que depenen de nuclis tensorials (Tensor cores), requereixen un treball explícit d'adaptació. No obstant això, gràcies a l'ús del DSL, els nuclis agnòstics d'arquitectura s'adapten de manera trivial.

Una GPU Blackwell té fins a 160 SM, en contrast amb els 132 de Hopper. Per a la majoria de nuclis, això és una constant en temps de compilació i la lògica del codi del llançador del costat de l'amfitrió ho té en compte per determinar el nucli òptim (preompliment vs descodificació) a executar. A més, Blackwell també introdueix instruccions de càrrega de 256 bits, que es poden emetre fàcilment augmentant l'amplada del vector (última dimensió dels tensors CuTe) per millorar la utilització de l'amplada de banda de la memòria.

Blackwell introdueix suport de maquinari per a un nombre significatiu d'esquemes de quantificació (MXFP8, MXFP4, NVFP4). El suport per a aquests es va implementar principalment ajustant els nuclis existents per a l'escalat per blocs de FP8. El suport JIT és crucial aquí, ja que es poden admetre nous modes sense fer explotar l'espai de configuració en temps de compilació. En el passat, vam trobar que Triton era incòmode a l'hora d'expressar reduccions horitzontals a través de pocs elements, sobre la meitat o un quart d'un warp. Un control més detallat sobre les reduccions era crucial per aconseguir un rendiment òptim en calcular un valor màxim absolut a través de 16 o 32 elements per a NVFP4 i MXFP8/MXFP4, respectivament.

Com a llenguatge relativament nou, els LLMs van tenir poca exposició al codi CuTeDSL durant el seu preentrenament. A més, l'estil del nostre codi actual, amb nuclis optimitzats per a activacions i normes, era lleugerament diferent dels exemples existents de CUTLASS o Quack, que estaven altament optimitzats per a TMA i instruccions de multiplicació de matrius específiques de l'arquitectura. Arribar a un punt on els LLMs poguessin ajudar en l'escriptura de nuclis va ser un procés gradual.

Vam construir manualment un corpus inicial de nuclis i funcions d'utilitat, seguint patrons consistents, abans de poder confiar en agents per a la conversió de nuclis. A mesura que el corpus avançava, podíem realitzar traduccions aproximades de CUDA a CuTeDSL, seguides d'un treball manual significatiu posterior. En particular, vam descobrir que els models tenien dificultats amb les conversions de punter a tensor i la vectorització, encallant-se en intentar arreglar errors de tipus quan l'objectiu era fer que una prova es compilés. Han mostrar una comprensió escassa de la naturalesa basada en graella dels nuclis, confiant la sincronització de fils i de graella. De manera més frustrant, hi havia la temptació d'al·lucinar traduccions inexistents de biblioteques específiques de CUDA en lloc de programar a partir dels nostres propis reemplaçaments de CuTeDSL. Amb el temps, a mesura que el corpus de nuclis CuTeDSL creixia, vam veure que els LLMs poden construir la major part del codi d'esquelet i omplir els espais en blanc, permetent-nos centrar-nos en el disseny d'alt nivell.

En comparació amb CUDA, trobem que CuTeDSL ofereix un rendiment en el temps de compilació molt millor, permetent-nos especialitzar-nos per a més paràmetres i aconseguir millores de rendiment. Tot i que els nuclis són lleugerament més verbosos a causa de la manca de determinades abstraccions en comparació amb Triton, considerem que la complexitat està justificada, ja que un control més granular sobre el maquinari ens permet construir nuclis més ràpids.

En el moment d'escriure aquest article, la bretxa de rendiment entre els nuclis CUTLASS C++ i CuTeDSL és mínima, obrint oportunitats apassionants per compondre i fusionar nuclis GEMM d'alt rendiment amb comunicació entre dispositius, normalització o activacions. En apostar per aquest DSL, podrem extreure més rendiment dels nostres dispositius per millorar el rendiment global (throughput) a totes les nostres càrregues de treball, des de representacions vectorials fins a LLMs.

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

Inferència a Perplexity AI

Per què CuTeDSL?

Experiència de desenvolupament

Gran espai de configuració

Preomplir i descodificar l'especialització

Sincronització de graella (grid)

Nuclis CuTeDSL

Norma QK (QK Norm)

Despatx/Combinació de barreja d'experts (Mixture-of-Experts Dispatch/Combine)

Local

Node únic (Singlenode)

Paral·lelisme de dades (Data Parallel)

Peer-to-Peer

Suport per a Blackwell

Generació de nuclis assistida per IA

Conclusió i treball futur

Referències

Citació