CuTeDSL en Perplexity

En Perplexity, empregamos o motor de servizo optimizado para o tempo de execución (ROSE, polas súas siglas en inglés) de elaboración propia para dar servizo a modelos que van desde incrustacións ata LLM de billóns de parámetros. Situado detrás de todas as API, incluíndo Sonar, Search e Embeddings, o noso motor de execución creouse para capacitar aos investigadores para explorar os modelos e as arquitecturas máis axeitados, á vez que desprega rapidamente os mellores candidatos en produción. Para acadar isto, necesitamos ser quen de implementar, despregar e optimizar a inferencia para arquitecturas de última xeración a un ritmo rápido. Nesta entrada do blog, exploramos como os núcleos CuTeDSL integrados no noso motor de inferencia nos permiten construír os núcleos de GPU necesarios para levar os modelos ao seu rendemento máximo nas GPU de NVIDIA.

Os modelos personalizados, construídos por varios equipos en Perplexity, alóxanse internamente en GPU NVIDIA Hopper e Blackwell mediante o motor de inferencia ROSE construído polo equipo de Inferencia de IA. ROSE creouse inicialmente para dar un mellor servizo aos modelos Llama personalizados tanto para a descodificación de linguaxe como para a clasificación, cunha interface compatible co servidor de inferencia NVIDIA Triton. Co tempo, converteuse nun motor de inferencia completo que ofrece LLM substanciais, así como modelos de clasificación, puntuación e outros baseados en transformers para unha ampla gama de equipos para impulsar os produtos de Perplexity. O obxectivo de deseño de ROSE é ser sinxelo: co tempo, invístese esforzo en acadar o rendemento máximo, pero a prioridade principal é adaptarse ás arquitecturas e aos casos de uso en constante evolución para os nosos modelos.

O espazo de configuración do kernel entre tipos de datos, dimensións de cabeceira, constantes en tempo de compilación e axustes de lanzamento.

A un nivel alto, ROSE consiste nun motor que adapta un modelo a unha interface coa que os clientes poden interactuar. Os motores son altamente reutilizables, xa que xestionan a programación e a agrupación (batching) de solicitudes e recuperan tokens descodificados, incrustacións ou puntuacións dos modelos subxacentes. A inicialización de dispositivos, a canles de comunicación entre procesos e entre nós, e a carga de pesos son xestionadas por eles. Para os LLM, o motor encárgase da agrupación, a fragmentación e a mostraxe. Tamén é responsable da asignación de almacenamento KV, tanto para a atención completa como lineal, así como da coincidencia de prefixos. Os motores de incrustación son máis sinxelos, xa que só realizan agrupacións en liña. A execución está optimizada de tal xeito que o motor prepara a seguinte rolda de traballo mentres o modelo se está a executar no acelerador en paralelo.

ROSE expón unha rica colección de capas personalizadas para a modelaxe, que envolven núcleos de GPU construídos con CuTeDSL, Triton, CUDA, CUTLASS e cuBLAS. Para as implementacións de multiplicación de matrices e atención, confiamos principalmente nos núcleos de NVIDIA de CUTLASS e cuBLAS. O resto da canalización, incluídos os núcleos de incrustación, norma, MoE e activación, forman parte da biblioteca de núcleos de ROSE. Para admitir mellor o hardware no que executamos a inferencia, ROSE está equipado coas súas propias primitivas de comunicación entre dispositivos para o enrutamento, envío e combinación de MoE. Aínda que os outros núcleos son relativamente sinxelos, deben especializarse a fondo para acadar o rendemento máximo en todas as diferentes dimensións ocultas e dispositivos nos que se executan. Ademais, confiamos na especialización en tempo de compilación para acomodar todos os botóns e axustes de varios modelos (sesgo de peso fronte a ningún sesgo de peso para a norma RMS, funcións de activación, etc.). En consecuencia, decidimos usar DSL para implementalos: inicialmente construímos ROSE usando Triton. Non entanto, desde entón adoptamos CuTeDSL como o contorno principal para a programación de GPU.

Desde unha perspectiva matemática, os núcleos de inferencia adoitan ser pouco sofisticados. En consecuencia, ao desenvolver un núcleo, a complexidade xeralmente non radica en expresar que cálculo se debe realizar, senón en como exactamente o hardware debe facelo. Para construír os mellores núcleos, necesitamos unha linguaxe que ofreza as abstraccións correctas sobre o hardware subxacente. A linguaxe correcta permítennos especificar as operacións de alto nivel para realizar o cálculo, especializándoos e optimizándoos de forma agresiva en todos os puntos dun espazo de configuración amplo.

CuTeDSL é unha DSL baseada en Python para GPU de NVIDIA construída sobre a álxebra de deseño CuTe e MLIR. Compílase xusto a tempo (JIT) en PTX altamente optimizado, mantendo ao mesmo tempo o control sobre as primitivas de hardware de baixo nivel. O seu obxectivo é ofrecer as mesmas características e rendemento que CUTLASS C++, sen os custosos tempos de compilación que dificultan a velocidade de desenvolvemento. Das linguaxes que examinamos ou coas que traballamos (incluíndo CUDA, cuTile, TVM IR e Triton), consideramos que é a mellor opción para o desenvolvemento dos nosos núcleos.

En xeral, aínda que non é tan de alto nivel e concisa como Triton ou cuTile, descubrimos que a álxebra de deseño CuTe e as súas composicións (expostas a través de cute.Tensor) eran substitutos razoables. Ademais, CuTeDSL ofrece acceso sen restricións ás primitivas de hardware, proporcionándonos o nivel de control que necesitamos para acadar o rendemento máximo en todos os casos límite. Con CUDA, as oportunidades de especialización en tempo de compilación eran limitadas, xa que cada parámetro adicional engadía un custo desrazonable aos tempos de compilación ou complexidade ao sistema de compilación. Pola contra, cunha DSL compilada xusto a tempo, podemos especializarnos de forma máis agresiva sen sacrificar o ritmo de iteración. A diferenza de TVM IR, que usaba CUDA como paso intermedio, descubrimos que a canalización baseada en MLIR que se reduciu a PTX era máis robusta e máis fácil de depurar que unha que emitía CUDA como paso intermedio. Ademais, co uso de operacións personalizadas para envolver as instrucións PTX, descubrimos que podíamos aproveitar ao máximo todas as características do hardware subxacente.

A maior parte do custo de compilación dos núcleos CUDA gasta normalmente na fase de expansión de modelos (templates) dos núcleos, xa que a optimización e a compilación desde a etapa PTX en diante pódense completar en cuestión de segundos. Con moitos parámetros en tempo de compilación, esta etapa de compilación pode explotar a varios minutos para un núcleo non trivial. Este cuello de botella acentuase polo feito de que as oportunidades de compilación paralela xorden só máis tarde, nos niveis IR e PTX, coa xestión de modelos de C++ que se produce de forma secuencial. Aínda que existen trucos para solucionar isto, ao mover as instanciacións de modelos a unidades de compilación separadas, requiren axustes significativos nos sistemas de compilación e o uso de macros inútiles. Pola contra, os núcleos CuTeDSL que se compilan Xusto a Tempo (JIT) non requiren regras de compilación complexas e non dificultan os tempos de compilación, xa que o ciclo de desenvolvemento pode centrarse nun único punto no espazo de configuración, mentres que as probas unitarias coidadosamente posteriores poden abranguer todo o espazo de deseño do núcleo.

Aínda que a compilación JIT move os custos do tempo de desenvolvemento ao tempo de execución, a sobrecarga é máis que aceptable. Cando se desprega un modelo, activa só algúns dos puntos de todo o espazo de configuración, xerando sobrecargas a escala de segundos durante o inicio/quecemento do motor de inferencia. Pola contra, a precompilación de todas as configuracións dun núcleo podería abranguer varios minutos. Ademais, as sobrecargas pódense mitigar aínda máis gardando en caché os artefactos compilados JIT.

O pesado custo de compilación de CUDA acentuase pola dificultade de rastrexar a orixe dos erros cando se utiliza o envío e a especialización baseados en modelos, como é o caso de CuTe e CUTLASS. Cun DSL, os erros pódense sinalar á liña de código que os xerou, facilitando significativamente a depuración. Ademais, é moito máis sinxelo mapear o código xerado coa sentenza DSL que o orixinou a través de etapas intermedias de MLIR ou confiando na información de depuración asociada ao PTX xerado. Grazas a estes aspectos da DSL, atopamos que somos moito máis produtivos ao desenvolver núcleos.

Aínda que os núcleos comparten moitas similitudes entre modelos, aínda deben estar equipados cunha serie de botóns de configuración que poden ter un impacto no rendemento. Os parámetros relevantes inclúen:

Dimensión oculta / Dimensión de cabeza / Número de cabezas de atención

Tipos de datos de entrada/saída

Esquema de cuantificación

Constantes en tempo de compilación (sesgos, funcións de activación, etc.)

Arquitectura ROSE, incluíndo o seu motor, modelo, biblioteca de kernel e servizos.

A especialización en tempo de compilación é útil tanto para o rendemento como para a depuración. Particularmente nun procesador en canalización (pipelined), como é o caso das GPU, unha rama inexistente é mellor que unha rama sempre tomada ou nunca tomada. Ademais, compilar un núcleo para un conxunto específico de botóns de configuración dá como resultado menos instrucións PTX/SASS, que son máis fáciles de analizar e depurar cando se usa cuda-gdb durante unha investigación.

Para un tamaño de lote e unha configuración determinados, pódense estimar o tráfico previsto ao DRAM e o número de operacións aritméticas. Posteriormente, pódense tomar decisións óptimas caso por caso para a asignación óptima da grella (grid), a sincronización e o uso de certas características de hardware. CuTeDSL permítenos definir e reutilizar compoñentes para unha especialización profunda sen incorrer nunha sobrecarga inaceptable en tempo de compilación.

En xeral, o rendemento dos núcleos de inferencia debe optimizarse para dous casos de uso separados:

Preenchido (Prefill): tamaño de lote pequeno (poucos fragmentos, secuencias), moitos tokens por secuencia. O rendemento debe maximizarse tanto para as incrustacións como para os tokens de entrada dos LLM. A latencia está entre decenas e centos de milisegundos.

Descodificación (Decode): tamaños de lote pequenos a grandes, dun (sen descodificación especulativa) a algúns tokens por secuencia (especulación de obxectivo de borrador, predición de varios tokens). A latencia está preto da latencia de lanzamento do núcleo.

Para GEMM e atención, a etapa de preenchido execútase nun gran número de tokens e o rendemento está limitado polo rendemento de cálculo do dispositivo. A descodificación implica a mesma operación en tamaños de lote pequenos, cos tempos de execución dominados polo custo de cargar as activacións e os pesos desde a DRAM nos multiprocesadores de fluxo (streaming multiprocessors) e núcleos de tensor, limitados pola ancho de banda de memoria do dispositivo no seu lugar. Aínda que o tamaño do lote ou o número total de tokens nun lote é un indicador relativamente claro, o limiar de cambio entre o preenchido e a descodificación é máis subtil porque o ancho de banda de memoria tamén está determinado pola dimensión oculta do modelo e o tipo de datos da activación. En consecuencia, os núcleos divídense en implementacións de preenchido e descodificación se as diferenzas algorítmicas ou as asignacións de grella/bloque beneficiarían a calquera deles.

O parámetro que é universalmente diferente entre os núcleos de preenchido e descodificación é o número de warps nun bloque (o número de fíos, dividido polo tamaño de warp de 32 fíos por warp). É desexable facer deste parámetro unha constante en tempo de compilación, xa que pode axudar na asignación de rexistros. En CUDA, isto faise normalmente especificando o número de fíos a través de __launch_bounds__ e fornecendo o número de warps como parámetro de modelo:

cpp

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

}

O número de rexistros dispoñibles para todos os fíos nun bloque é unha restrición de hardware, 64K en Hopper e Blackwell. Se o número de fíos se coñece de antemán, o compilador sabe exactamente cantos rexistros pode asignar a cada fío e pode ter éxito ao compilar núcleos optimizados de forma máis agresiva. Non entanto, isto ten un custo en tempo de compilación, xa que agora o núcleo ten que ser recompilado para diferentes tamaños de warp. Con CuTeDSL, isto xa non é un problema, xa que non necesitamos pagar o custo da compilación para todo o espazo de configuración por adiante. A DSL tamén agarda que o tamaño do bloque sexa unha expresión constante, realizando implicitamente a optimización implícita por __latch_bounds__ en CUDA.

Ademais de facer que o número de warps sexa unha constante, é desexable varialo entre o preenchido e a descodificación. Para utilizar plenamente o ancho de banda de memoria dunha GPU, non é suficiente asignar sempre un alto número de warps por bloque, xa que o ancho de banda completo só se pode conseguir usando tantos multiprocesadores de fluxo dispoñibles (SM) como sexa posible. En consecuencia, para un número determinado de tokens e tipos de datos e dimensións ocultas coñecidos, primeiro escollemos unha dimensión vectorial ideal por fío (8 para bfloat16, 4 para float32) para garantir que o código xerado use as instrucións de carga/almacenamento de 128 bits (256 bits en Blackwell). Coñecendo o tamaño do vector, podemos calcular o número de warps que require o núcleo. Primeiro intentamos maximizar o número de SM: se hai menos warps que SM, elíxese un único warp por bloque. Esta adoita ser a variante de descodificación do núcleo. Se hai suficientes warps, aumentamos gradualmente o número de warps a 8, 16 ou 32, avaliando un límite superior despois do cal normalmente xa non vemos unha mellora no rendemento.

Aínda que Triton realizou algunhas destas optimizacións baixo o capó, descubrimos que podemos acadar resultados similares utilizando coidadosamente a álxebra de deseño CuTe, sen estar limitados polas deficiencias dun compilador de alto nivel. Particularmente cando se producen reducións horizontais sub-warp, descubrimos que CuTeDSL nos ofrecía o control de gran fino necesario para xerar código óptimo.

Nalgunhas situacións, para distribuír o traballo entre tantos bloques como sexa posible, pódese usar cooperative_groups::this_grid().sync() para sincronizar todos os bloques participantes. Aplícanse restricións: dado que un bloque non se pode desprogramar dun SM unha vez que comeza a execución, para garantir o progreso cara adiante, o número de bloques non pode superar o número de SM dispoñibles no dispositivo, para garantir que en calquera punto todos os bloques poidan actualizar e sondar as localizacións de memoria global implicadas na sincronización. Esta restricción é imposta por cudaLaunchCooperativeKernel. O custo dunha barreira de grella é de aproximadamente 3 us en Hopper, máis o custo de baleirar todas as operacións de memoria anteriores a través dunha barreira de memoria implícita ngunhas situacións. En xeral, isto é comparable á sobrecarga de dividir un núcleo en dous lanzamentos, o que leva a unha decisión de deseño sensible ao rendemento. A pesar de que o lanzamento dependente programático (PDL) pode reducir as sobrecargas de lanzamento desde Hopper en diante, para os núcleos de descodificación, vale a pena evitar a sobrecarga.

Se un dos cargas de traballo require máis bloques que o número de SM, o núcleo debe dividirse, confiando na barreira implícita dos lanzamentos de núcleos. Este é principalmente o caso do preenchido. Se o número de bloques que son ideais para todas as cargas de traballo encaixa na grella, prefírese unha barreira de grella, como é o caso dos núcleos de descodificación que operan nun número menor de tokens.

En CuTeDSL, estas decisións pódense tomar no código de lanzamento do lado do host e os núcleos pódense compoñer e entretecer con barreiras de grella segundo sexa necesario, como se ilustra no exemplo seguinte. Con tal estratexia híbrida, os núcleos de descodificación pódense optimizar para acadar latencias mínimas e os de preenchido pódense optimizar para un rendemento máximo. A diferenza das solucións baseadas en CUDA, esta lóxica pódese definir de forma máis concisa e con menos sobrecarga en tempo de compilación.

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

No caso dos nosos núcleos all-to-all de MoE, como se detalla a continuación, ter a flexibilidade de construír núcleos con ou sen barreiras de grella e optimizalos en tempo de compilación axudounos a acadar ganancias de rendemento substanciais a través de CuTeDSL.

O primeiro núcleo que construímos en CuTeDSL foi un núcleo de norma RMS especializado, optimizado para dimensións pequenas. A norma QK, utilizada na capa de auto-atención dos modelos Qwen recentes, aplícase por cabeza, en 64, 128 ou 256 elementos. A diferenza da norma RMS normal, que opera en dimensións ocultas que superan os milhares, o que require unha redución horizontal tanto a través da memoria compartida como dos rexistros, a norma QK necesita redución só a través de warps, coa posibilidade de comprimir múltiples cabezas nun único warp. Isto converteuse nun caso de uso no que podíamos optimizar agresivamente en tempo de compilación para atopar as mellores asignacións de token a warp para todas as dimensións de cabeza e tipos de datos.

A norma RMS é unha operación bastante directa, aínda que require unha suma horizontal para calcular a norma ao cadrado dun vector:

cpp

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

Dado que as cabezas de atención adoitan ser pequenas e as operacións por fío están moi vectorizadas, un único warp pode calcular a norma ao cadrado. Co fin de maximizar o ancho de banda de memoria, cada fío debe procesar 128 bits de entrada, o que se traduce en 8 valores de bfloat16. En consecuencia, un medio warp é suficiente para procesar unha cabeza, e a norma ao cadrado pódese determinar mediante unha redución horizontal realizada por shuffle_sync_bfly. Partindo da dimensión da cabeza e do tipo de datos, podemos determinar o número de cabezas por warp e especificar os argumentos e o número de roldas de redución no momento da compilación.

Despois podemos construír un núcleo especializado con relativamente poucas liñas de código, confiando en cargas/almacenamentos vectorizados automaticamente e usando operacións específicas de arquitectura de baixo nivel para a redución de warp. O núcleo determina en que cabeza (Q ou K) traballar e identifica os pesos correspondentes para aplicar. Carga un vector completo da cabeza e primeiro suma localmente. Posteriormente, en función do número de fíos por cabeza, calcúlase a norma ao cadrado. Finalmente, a norma RMS, vectorizada automaticamente, determínase e escríbese de novo na memoria. Os parámetros específicos do modelo, como WEIGHT_BIAS e EPS, están plegados en constantes (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ón cun núcleo de norma RMS normal, esta variante especializada ofrece un rendemento e unha latencia de 2 a 3 veces mellores tanto para o preenchido como para a descodificación.

O envío e a combinación de MoE barallan os tokens de entrada aos expertos aos que se enrutan, realizando unha acumulación ponderada despois de que os núcleos GEMM de MoE realicen as multiplicacións de matrices relevantes. O rendemento é de especial interese, xa que varios produtos confían en modelos MoE de varios tamaños. Dependendo dos tipos de descargas, distinguimos diferentes núcleos para optimizar:

TP=1: Núcleo local para unha única GPU. O envío baralla os tokens aos expertos relevantes e a combinación acúmulasos. Utilizado principalmente para despregamentos de Qwen3 30B en H100.

DP=1, TP ≠ 1: Núcleo tensor-paralelo, dun só nodo. A entrada replícase, o envío selecciona os tokens para o rango actual e a combinación acúmulaos en todos os rangos. Acumulación a través de NVLink ou NVLS. Utilizado para preenchedores grandes.

DP ≠ 1, NVLink: Núcleo de datos paralelos con paralelismo de tensores opcional. Cada rango envía tokens a cada outro rango que aloxa expertos específicos. A información de enrutamento agrúpase, os tokens transfirense a través de NVIDIA NVLink. Utilizado para despregamentos H200 ou GB200 paralelos aos datos.

DP ≠ 1, InfiniBand: Núcleo peer-to-peer paralelo aos datos e paralelo aos tensores que opera a través de NVLink e NVIDIA Quantum InfiniBand. Utilizado para despregamentos Hopper multinodo.

O núcleo local conta o número de tokens asignados a cada experto en función da táboa de enrutamento e asigna un índice único a cada token dentro da lista de tokens que pertencen a cada experto. Despois do paso de enrutamento, os tokens cópianse en búferes opcionalmente recheos (padded) para pasar a un núcleo GEMM de MoE. Utilizando as localizacións dos tokens, o núcleo de combinación realiza entón unha acumulación nos rexistros, escribindo os tokens finais nun búfer de saída.

A fase de envío varía entre o preenchido e a descodificación: para poucos tokens, hai pouca información de enrutamento para ler, polo que a procesamos nun único bloque e sinalamos o resto da grella, con bloques individuais copiando fragmentos de tokens. A asignación de tokens determínase incrementando atomicamente os contadores na memoria compartida. Dado que o preenchido debe procesar máis información de enrutamento, o enrutamento divídese en múltiples bloques nun núcleo separado. Os índices asígnanse primeiro localmente usando contadores na memoria compartida e despois agrúpanse en todos os bloques na memoria global para determinar a posición de cada token. A grella deste núcleo bótase (clamped) ao número de SM para poder usar barreiras a nivel de grella. Un núcleo posterior realiza as copias, lanzado nun número moito maior de bloques e warps para utilizar plenamente o ancho de banda de memoria dispoñible.

O núcleo de combinación é relativamente sinxelo, xa que realiza unha acumulación vectorizada, baseándose na información de enrutamento e nos desprazamentos (offsets) calculados polo envío. CuTeDSL demostrou ser valioso aquí, xa que o cálculo podería expresarse a través de tensores CuTe simples que compilan en código eficiente para todos os tipos de datos e dimensións ocultas.

Nun despregamento tensor-paralelo, a información de enrutamento e os tokens de entrada replícanse. O núcleo de envío local pódese reutilizar completamente cun lixeiro cambio para enrutar e seleccionar só os tokens asignados aos expertos situados no rango actual. Con algúns parámetros e especialización en tempo de compilación, os núcleos de envío local reutilízanse.

Dado que o número de rangos adoita ser menor que o número de expertos aos que se enruta un token, no caso medio pódese asumir que cada token ten polo menos unha das súas réplicas enrutadas a cada rango. Nesta situación, o núcleo de combinación local pódese reutilizar para agregar as copias locais dun token, seguido dun all-reduce para sumar en todos os rangos participantes. No caso do núcleo de preenchido, lanzamos núcleos separados, porén a descodificación fusiona ambas operacións nun único lanzamento.

Para all-reduce, usamos NVLink simple se só participan dous rangos, cambiando a NVLS cando interveñen 4 ou máis rangos para un aumento do 10-20% no rendemento. Grazas á DSL, podemos compartir a maior parte da lóxica de redución e os intercambios de búferes de pares entre o MoE e os núcleos all-reduce regulares. Esta é unha situación na que o control estreito sobre o hardware é beneficioso, xa que se requiren operacións DSL sobre multimem.ld_reduce e multimem.st para NVLS. A implementación da redución en si segue en gran medida o exemplo de all_reduce das probas NCCL.

Os núcleos paralelos aos datos confían en NVLink para agregar tokens para múltiples dispositivos. na fase de envío, antes de intercambiar tokens, os pares primeiro enrutan localmente para identificar o desprazamento de destino en cada rango para cada token que están enviando. Posteriormente, comezan a copiar os tokens en búferes de recepción asignados para cada emisor para evitar conflitos. As escrituras de liberación-almacenamento (release-store) por token utilízanse para sinalar a finalización. No extremo do receptor, a información de enrutamento agrúpase para poder condensar todos os tokens nun único búfer para o GEMM de MoE. Os receptores sondan a bandeira escrita polo par cunha carga de adquisición (acquire-load) para esperar a que rematen as escrituras.

No lado da combinación, os tokens cópianse a través de NVLink aos rangos de orixe, cunha bandeira por token que indica a recepción. O núcleo de combinación é semellante ás variantes local e dun só nodo, con lixeiras diferenzas no enderezo e algunha lóxica de sincronización adicional.

Construímos unha biblioteca de contenedores (wrappers) PTX para expoñer cargas e almacenamentos de liberación-adquisición con ámbito tanto a nivel de sistema como de dispositivo para implementar este núcleo. Ao mellorar a granularidade das bandeiras de sincronización e asignar máis bloques para enviar e recibir tokens a través de NVLink, conseguimos melloras de 2 veces sobre os nosos núcleos CUDA anteriores.

Os nosos núcleos Peer-to-Peer conseguiron anteriormente un rendemento de última xeración en GPU Hopper combinadas con adaptadores InfiniBand ConnectX-7. Desde entón portamos os núcleos de combinación a CuTeDSL para obter máis ganancias de rendemento. Dado que as versións de CUDA xa confiaban en gran medida na especialización en tempo de compilación, demostraron ser os cuellos de botella en tempo de compilación no noso repositorio, obrigándonos a omitir algúns parámetros. Ao migrar á DSL compilada por JIT, especializámonos en máis parámetros e axustamos as asignacións de grella, introducindo máis diferenzas entre o preenchido e a descodificación para un impulso de arredor do 3% na latencia.

Desde a perspectiva dos núcleos ROSE, a compatibilidade con Blackwell divídese en dúas categorías: núcleos agnósticos de arquitectura e específicos de arquitectura. Os núcleos específicos de arquitectura, que confían nos núcleos Tensor, requiren un traballo explícito para adaptarse. Non entanto, grazas ao uso da DSL, os núcleos agnósticos de arquitectura adáptanse de forma trivial.

Unha GPU Blackwell ten ata 160 SM, en contraste cos 132 de Hopper. Para a maioría dos núcleos, esta é unha constante en tempo de compilación e a lóxica no código do lanzador do lado do host ténea en conta para determinar o núcleo óptimo (preenchido fronte a descodificación) para executar. Ademais, Blackwell tamén introduce instrucións de carga de 256 bits, que se poden emitir facilmente aumentando a ancho do vector (última dimensión dos tensores CuTe) para mellorar a utilización do ancho de banda de memoria.

Blackwell introduce compatibilidade de hardware para un número significativo de esquemas de cuantificación (MXFP8, MXFP4, NVFP4). A compatibilidade con estes implementouse principalmente axustando os núcleos existentes para o escalado por bloques FP8. A compatibilidade JIT é fundamental aquí, xa que se poden admitir novos modos sen que explote o espazo de configuración en tempo de compilación. No pasado, descubrimos que Triton era incómodo á hora de expresar reducións horizontais en poucos elementos, sobre a metade ou un cuarto dun warp. O control máis detallado sobre as reducións foi fundamental para acadar un rendemento óptimo ao calcular un valor máximo absoluto en 16 ou 32 elementos para NVFP4 e MXFP8/MXFP4, respectivamente.

Como linguaxe relativamente nova, os LLM tiveron poca exposición ao código CuTeDSL durante o seu preadestramento. Ademais, o estilo do noso código actual, con núcleos optimizados para activacións e normas, era lixeiramente diferente dos exemplos existentes de CUTLASS ou Quack, que estaban fortemente optimizados para TMA e instrucións de multiplicación de matrices específicas de arquitectura. Chegar a un punto no que os LLM puidesen axudar na escritura de núcleos foi un proceso gradual.

Construímos manualmente un corpus inicial de núcleos e funcións de utilidade, seguindo patróns consistentes, antes de poder confiar nos axentes para a conversión de núcleos. A medida que avanzaba o corpus, podíamos realizar traducións aproximadas de CUDA a CuTeDSL, cun traballo manual significativo despois. En particular, descubrimos que os modelos tiñan dificultades coas conversións de punteiro a tensor e a vectorización, quedando atrapados tentando arranxar erros de tipo cando o obxectivo era facer que unha proba compilase. Mostraron pouco entendemento da natureza baseada en grellas dos núcleos, confundindo a sincronización de fíos e grellas. O máis frustrante foi a tentación de alucinacións de traducións non existentes de bibliotecas específicas de CUDA en lugar de construír en contra dos nosos propios substitutos de CuTeDSL. Co tempo, a medida que crecía o corpus de núcleos CuTeDSL, descubrimos que os LLM poden construír a maior parte do código esqueleto e encher os espazos en branco, permitíndonos centrarnos no deseño de alto nivel.

En comparación con CUDA, atopamos que CuTeDSL ofrece un rendemento en tempo de compilación moito mellor, permitíndonos especializarnos en máis parámetros e acadar melloras de rendemento. Aínda que os núcleos son un pouco máis verbosos debido á falta de certas abstraccións en comparación con Triton, atopamos que a complexidade está xustificada xa que un control máis detallado sobre o hardware nos permite construír núcleos máis rápidos.

No momento de escribir este artigo, a brecha de rendemento entre os núcleos CUTLASS C++ e CuTeDSL é mínima, abrindo oportunidades emocionantes para compoñer e fusionar núcleos GEMM de alto rendemento coa comunicación entre dispositivos, normalización ou activacións. Ao comprometernos con esta DSL, poderemos extraer máis rendemento dos nosos dispositivos para mellorar o rendemento nas nosas cargas de traballo, desde incrustacións ata LLM.

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

Inferencia en Perplexity AI

Por que CuTeDSL?

Experiencia de desenvolvemento

Espazo de configuración grande

Preencher e decodificar a especialización

Sincronización de grellas

Núcleos CuTeDSL

Norma QK

Envío/Combinación de Mixture-of-Experts (MoE)

Local

Dun só nodo

Paralelo aos datos

Peer-to-Peer

Compatibilidade con Blackwell

Xeración de núcleos asistida por IA

Conclusión e traballo futuro

Referencias

Cita