CuTeDSL en Perplexity

En Perplexity, utilizamos el motor de servicio optimizado para tiempo de ejecución (ROSE, por sus siglas en inglés) de desarrollo propio para ofrecer modelos que van desde incrustaciones hasta LLMs de un billón de parámetros. Situado detrás de todas las API, incluidas Sonar, Search y Embeddings, nuestro motor de ejecución se diseñó para permitir a los investigadores explorar los modelos y arquitecturas más adecuados, al tiempo que implementan rápidamente los mejores candidatos en producción. Para lograr esto, necesitamos poder implementar, desplegar y optimizar la inferencia para arquitecturas de última generación a un ritmo acelerado. En esta publicación de blog, exploramos cómo los núcleos CuTeDSL integrados en nuestro motor de inferencia nos permiten construir los núcleos de GPU necesarios para llevar los modelos al máximo rendimiento en GPU de NVIDIA.

Los modelos personalizados, creados por varios equipos de Perplexity, se alojan internamente en GPU NVIDIA Hopper y Blackwell utilizando el motor de inferencia ROSE creado por el equipo de Inferencia de IA. ROSE se creó inicialmente para dar un mejor servicio a los modelos Llama personalizados tanto para la decodificación de lenguaje como para la clasificación, con una interfaz compatible con NVIDIA Triton Inference Server. Con el tiempo, se convirtió en un motor de inferencia completo que ofrece LLMs sustanciales, así como modelos de clasificación, puntuación y basados en transformadores para una amplia gama de equipos para impulsar los productos de Perplexity. El objetivo de diseño de ROSE es ser simple: con el tiempo, se invierte esfuerzo para alcanzar el rendimiento máximo, pero la prioridad principal es adaptarse a las arquitecturas y casos de uso en constante evolución para nuestros modelos.

El espacio de configuración del kernel en tipos de datos, dimensiones de cabezas, constantes en tiempo de compilación y configuraciones de inicio.

A un alto nivel, ROSE consta de un motor que adapta un modelo a una interfaz con la que los clientes pueden interactuar. Los motores son altamente reutilizables; gestionan la programación y agrupación de solicitudes, y recuperan tokens decodificados, incrustaciones o puntuaciones de los modelos subyacentes. La inicialización de dispositivos, los canales de comunicación entre procesos y entre nodos, y la carga de pesos son gestionados por ellos. Para los LLMs, el motor es responsable de la agrupación (batching), la fragmentación (chunking) y el muestreo. También es responsable de la asignación de almacenamiento KV, tanto para la atención completa como lineal, así como de la coincidencia de prefijos. Los motores de incrustación son más simples, ya que solo realizan agrupaciones en línea. La ejecución está optimizada de modo que el motor prepara la siguiente ronda de trabajo mientras el modelo se ejecuta en el acelerador en paralelo.

ROSE expone una rica colección de capas personalizadas para modelado, envolviendo núcleos de GPU construidos con CuTeDSL, Triton, CUDA, CUTLASS y cuBLAS. Para las implementaciones de multiplicación de matrices y atención, confiamos principalmente en los núcleos de NVIDIA de CUTLASS y cuBLAS. El resto de la canalización, incluidos los núcleos de incrustación, norma, MoE y activación, forman parte de la biblioteca de núcleos de ROSE. Para admitir mejor el hardware en el que ejecutamos la inferencia, ROSE está equipado con sus propias primitivas de comunicación entre dispositivos para el enrutamiento, envío y combinación de MoE. Si bien los otros núcleos son relativamente simples, deben estar altamente especializados para lograr el rendimiento máximo en todas las diferentes dimensiones ocultas y dispositivos en los que se ejecutan. Además, dependemos de la especialización en tiempo de compilación para acomodar todos los ajustes y modificaciones de varios modelos (sesgo de peso frente a sin sesgo de peso para la norma RMS, funciones de activación, etc.). En consecuencia, decidimos usar DSL para implementarlos: inicialmente construimos ROSE usando Triton. Sin embargo, desde entonces hemos adoptado CuTeDSL como el entorno principal para la programación de GPU.

Desde una perspectiva matemática, los núcleos de inferencia suelen ser poco sofisticados. En consecuencia, al desarrollar un núcleo, la complejidad generalmente no radica en expresar qué cálculo se debe realizar, sino en cómo exactamente debe hacerlo el hardware. Para construir los mejores núcleos, requerimos un lenguaje que ofrezca las abstracciones adecuadas sobre el hardware subyacente. El lenguaje correcto nos permite especificar las operaciones de alto nivel para realizar el cálculo, especializándolas y optimizándolas agresivamente en todos los puntos de un amplio espacio de configuración.

CuTeDSL es un DSL basado en Python para GPU de NVIDIA construido sobre el álgebra de diseño CuTe y MLIR. Se compila justo a tiempo (JIT) en PTX altamente optimizado, al tiempo que conserva el control sobre las primitivas de hardware de bajo nivel. Su objetivo es ofrecer las mismas características y rendimiento que CUTLASS C++, sin los costosos tiempos de compilación que obstaculizan la velocidad de desarrollo. De los lenguajes que hemos analizado o con los que hemos trabajado (incluyendo CUDA, cuTile, TVM IR y Triton), consideramos que es la mejor opción para el desarrollo de nuestros núcleos.

En general, aunque no es tan de alto nivel y conciso como Triton o cuTile, descubrimos que el álgebra de diseño CuTe y sus composiciones (expuestas a través de cute.Tensor) eran sustitutos razonables. Además, CuTeDSL ofrece acceso sin restricciones a las primitivas de hardware, proporcionándonos el nivel de control que requerimos para lograr el rendimiento máximo en todos los casos extremos. Con CUDA, las oportunidades de especialización en tiempo de compilación eran limitadas, ya que cada parámetro adicional añadía un costo irrazonable a los tiempos de compilación o complejidad al sistema de compilación. En contraste, con un DSL compilado justo a tiempo, podemos especializarnos de manera más agresiva sin sacrificar el ritmo de iteración. A diferencia de TVM IR, que usaba CUDA como un paso intermedio, descubrimos que la canalización basada en MLIR que descendía a PTX era más robusta y fácil de depurar que una que emitía CUDA como paso intermedio. Además, con el uso de operaciones personalizadas para envolver instrucciones PTX, descubrimos que podíamos aprovechar al máximo todas las características del hardware subyacente.

La mayor parte del costo de compilación de los núcleos CUDA se gasta típicamente en la fase de expansión de plantillas de los núcleos, ya que la optimización y la compilación desde la etapa de PTX en adelante se pueden completar en cuestión de segundos. Con muchos parámetros en tiempo de compilación, esta etapa de compilación puede dispararse a varios minutos para un núcleo no trivial. Este cuello de botella se ve acentuado por el hecho de que las oportunidades para la compilación en paralelo surgen solo más tarde, en los niveles de IR y PTX, y el manejo de plantillas de C++ ocurre secuencialmente. Aunque existen trucos para abordar esto, al mover las instanciaciones de plantillas a unidades de compilación separadas, requieren modificaciones significativas en los sistemas de compilación y el uso de macros difíciles de manejar. En contraste, los núcleos CuTeDSL que se compilan justo a tiempo (JIT) no requieren reglas de compilación complejas ni obstaculizan los tiempos de compilación, ya que el ciclo de desarrollo puede centrarse en un solo punto en el espacio de configuración, mientras que pruebas unitarias cuidadosas posteriores pueden cubrir todo el espacio de diseño del núcleo.

Aunque la compilación JIT traslada los costos del tiempo de desarrollo al tiempo de ejecución, la sobrecarga es más que aceptable. Cuando se despliega un modelo, este activa solo unos pocos puntos de todo el espacio de configuración, lo que genera sobrecargas de escala de segundos durante el inicio y calentamiento del motor de inferencia. En contraste, precompilar todas las configuraciones de un núcleo podría abarcar varios minutos. Además, las sobrecargas se pueden mitigar aún más almacenando en caché los artefactos compilados por JIT.

El fuerte costo de compilación de CUDA se ve acentuado por la dificultad de rastrear el origen de los errores cuando se utiliza el envío y la especialización basados en plantillas, como es el caso de CuTe y CUTLASS. Con un DSL, los errores se pueden señalar en la línea de código que los generó, lo que facilita significativamente la depuración. Además, es mucho más simple mapear el código generado a la instrucción DSL que lo originó a través de etapas intermedias de MLIR o basándose en la información de depuración adjunta al PTX generado. Gracias a estos aspectos del DSL, nos encontramos mucho más productivos al desarrollar núcleos.

Si bien los núcleos comparten muchas similitudes entre modelos, aún deben estar equipados con una serie de opciones de configuración que pueden tener un impacto en el rendimiento. Los parámetros relevantes incluyen:

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

Tipos de datos de entrada/salida

Esquema de cuantización

Constantes en tiempo de compilación (sesgos, funciones de activación, etc.)

Arquitectura ROSE, incluyendo su motor, modelo, biblioteca del kernel y servicios.

La especialización en tiempo de compilación es útil tanto para el rendimiento como para la depuración. Particularmente en un procesador canalizado, como es el caso de las GPU, una rama inexistente es mejor que una rama que siempre se toma o nunca se toma. Además, compilar un núcleo para un conjunto específico de opciones de configuración da como resultado menos instrucciones PTX/SASS, que son más fáciles de analizar y depurar al usar cuda-gdb durante una investigación.

Para un tamaño de lote y configuración dados, se puede estimar el tráfico esperado hacia la DRAM y el número de operaciones aritméticas. Posteriormente, se pueden tomar decisiones óptimas caso por caso para la asignación óptima de la cuadrícula, la sincronización y el uso de ciertas características de hardware. CuTeDSL nos permite definir y reutilizar componentes para una especialización profunda sin incurrir en una sobrecarga inaceptable en el tiempo de compilación.

En general, el rendimiento de los núcleos de inferencia debe optimizarse para dos casos de uso separados:

Prellenado (Prefill): tamaño de lote pequeño (pocos fragmentos, secuencias), muchos tokens por secuencia. Rendimiento a maximizar tanto para incrustaciones como para los tokens de entrada de los LLMs. Latencia de decenas a cientos de milisegundos.

Decodificación (Decode): tamaños de lote de pequeños a grandes, de uno (sin decodificación especulativa) a unos pocos tokens por secuencia (especulación de borrador a destino, predicción de múltiples tokens). Latencia cercana a la latencia de lanzamiento del núcleo.

Para GEMM y la atención, la etapa de prellenado se ejecuta en una gran cantidad de tokens y el rendimiento está limitado por el rendimiento de cómputo del dispositivo. La decodificación implica la misma operación en tamaños de lote pequeños, con tiempos de ejecución dominados por el costo de cargar las activaciones y los pesos desde la DRAM en los multiprocesadores de transmisión y los núcleos tensoriales, limitados en su lugar por el ancho de banda de memoria del dispositivo. Aunque el tamaño del lote o el número total de tokens en un lote es un indicador relativamente claro, el umbral de cambio entre el prellenado y la decodificación es más sutil porque el ancho de banda de memoria también está determinado por la dimensión oculta del modelo y el tipo de datos de la activación. En consecuencia, los núcleos se dividen en implementaciones de prellenado y decodificación si las diferencias algorítmicas o las asignaciones de cuadrícula/bloque benefician a cualquiera de ellos.

El parámetro que es universalmente diferente entre los núcleos de prellenado y decodificación es el número de warps en un bloque (el número de hilos, dividido por el tamaño de warp de 32 hilos por warp). Es deseable hacer de este parámetro una constante en tiempo de compilación, ya que puede ayudar en la asignación de registros. En CUDA, esto se hace típicamente especificando el número de hilos a través de __launch_bounds__ y proporcionando el número de warps como parámetro de plantilla:

cpp

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

}

El número de registros disponibles para todos los hilos en un bloque es una restricción de hardware, de 64K en Hopper y Blackwell. Si se conoce el número de hilos de antemano, el compilador sabe exactamente cuántos registros puede asignar a cada hilo y puede tener éxito en la compilación de núcleos optimizados de manera más agresiva. Sin embargo, esto tiene un costo en el tiempo de compilación, ya que ahora el núcleo debe recompilarse para diferentes tamaños de warp. Con CuTeDSL, esto ya no es un problema, ya que no necesitamos pagar el costo de compilación para todo el espacio de configuración por adelantado. El DSL también espera que el tamaño del bloque sea una expresión constante, realizando implícitamente la optimización implicada por __latch_bounds__ en CUDA.

Además de hacer que el número de warps sea una constante, es deseable variarlo entre el prellenado y la decodificación. Para utilizar plenamente el ancho de banda de memoria de una GPU, no es suficiente asignar siempre un alto número de warps por bloque, ya que el ancho de banda completo solo se puede lograr utilizando tantos Multiprocesadores de Transmisión (SM) disponibles como sea posible. En consecuencia, para un número dado de tokens y tipos de datos y dimensiones ocultas conocidos, primero elegimos una dimensión de vector ideal por hilo (8 para bfloat16, 4 para float32) para garantizar que el código generado utilice las instrucciones de carga/almacenamiento de 128 bits (256 bits en Blackwell). Conociendo el tamaño del vector, podemos calcular el número de warps requeridos por el núcleo. Primero intentamos maximizar el número de SM: si hay menos warps que SM, se elige un solo warp por bloque. Esta es típicamente la variante de decodificación del núcleo. Si hay suficientes warps, aumentamos gradualmente el número de warps a 8, 16 o 32, evaluando un límite superior después del cual normalmente ya no vemos una mejora en el rendimiento.

Si bien Triton realizó algunas de estas optimizaciones bajo el capó, descubrimos que podemos lograr resultados similares utilizando cuidadosamente el álgebra de diseño CuTe, sin estar limitados por las deficiencias de un compilador de alto nivel. Particularmente cuando se trata de reducciones horizontales de sub-warp, descubrimos que CuTeDSL nos ofrecía el control detallado requerido para generar código óptimo.

En algunas situaciones, para distribuir el trabajo entre tantos bloques como sea posible, se puede utilizar cooperative_groups::this_grid().sync() para sincronizar todos los bloques participantes. Se aplican restricciones: dado que un bloque no se puede desprogramar de un SM una vez que comienza la ejecución, para garantizar el progreso hacia adelante, el número de bloques no puede exceder el número de SM disponibles en el dispositivo, para garantizar que en cualquier punto todos los bloques puedan actualizar y sondear las ubicaciones de memoria global involucradas en la sincronización. Esta restricción es aplicada por cudaLaunchCooperativeKernel. El costo de una barrera de cuadrícula es de aproximadamente 3 us en Hopper, más el costo de vaciar todas las operaciones de memoria anteriores a través de una barrera de memoria implícita en algunas situaciones. En general, esto es comparable a la sobrecarga de dividir un núcleo en dos lanzamientos, lo que lleva a una decisión de diseño sensible al rendimiento. A pesar de que el lanzamiento dependiente programático (PDL) puede reducir las sobrecargas de lanzamiento desde Hopper en adelante, para los núcleos de decodificación, evitar la sobrecarga vale la pena.

Si una de las cargas de trabajo requiere más bloques que el número de SM, el núcleo debe dividirse, confiando en la barrera implícita de los lanzamientos de núcleos. Este es principalmente el caso del prellenado. Si el número de bloques que son ideales para todas las cargas de trabajo cabe en la cuadrícula, se prefiere una barrera de cuadrícula, como es el caso de los núcleos de decodificación que operan en un número menor de tokens.

En CuTeDSL, estas decisiones se pueden tomar en el código de lanzamiento del lado del host y los núcleos se pueden componer e intercalar con barreras de cuadrícula según sea necesario, como lo ilustra el ejemplo a continuación. Con una estrategia híbrida de este tipo, los núcleos de decodificación se pueden optimizar para alcanzar latencias mínimas y los de prellenado se pueden optimizar para un rendimiento máximo. A diferencia de las soluciones basadas en CUDA, esta lógica se puede definir de manera más concisa y con menos sobrecarga en el tiempo 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)
		)

En el caso de nuestros núcleos MoE all-to-all, como se detalla a continuación, tener la flexibilidad de construir núcleos con o sin barreras de cuadrícula y optimizarlos en el tiempo de compilación nos ayudó a lograr ganancias de rendimiento sustanciales a través de CuTeDSL.

El primer núcleo que construimos en CuTeDSL fue un núcleo de norma RMS especializado, optimizado para dimensiones pequeñas. La norma QK, utilizada en la capa de autoatención de los modelos Qwen recientes, se aplica por cabeza, en 64, 128 o 256 elementos. A diferencia de la norma RMS regular, que opera en dimensiones ocultas que superan los miles, requiriendo una reducción horizontal a través de la memoria compartida y los registros, la norma QK necesita reducción solo a través de los warps, con la posibilidad de exprimir múltiples cabezas en un solo warp. Esto se convirtió en un caso de uso en el que podíamos optimizar agresivamente en el tiempo de compilación para encontrar las mejores asignaciones de token a warp para todas las dimensiones de cabeza y tipos de datos.

La norma RMS es una operación bastante sencilla, aunque requiere una suma horizontal para calcular la norma al cuadrado de un vector:

cpp

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

Dado que las cabezas de atención suelen ser pequeñas y las operaciones por hilo están fuertemente vectorizadas, un solo warp puede calcular la norma al cuadrado. Para maximizar el ancho de banda de memoria, cada hilo debe procesar 128 bits de entrada, lo que se traduce en 8 valores de bfloat16. En consecuencia, un medio warp es suficiente para procesar una cabeza, y la norma al cuadrado se puede determinar utilizando una reducción horizontal realizada por shuffle_sync_bfly. A partir de la dimensión de la cabeza y el tipo de datos, podemos determinar el número de cabezas por warp y especificar los argumentos y el número de rondas de reducción en el momento de la compilación.

Luego podemos construir un núcleo especializado con relativamente pocas líneas de código, basándonos en cargas/almacenamientos vectorizados automáticamente y utilizando operaciones de nivel inferior específicas de la arquitectura para la reducción del warp. El núcleo determina en qué cabeza (Q o K) trabajar e identifica los pesos correspondientes a aplicar. Carga un vector completo de la cabeza y primero suma localmente. Posteriormente, según el número de hilos por cabeza, se calcula la norma al cuadrado. Finalmente, la norma RMS, vectorizada automáticamente, se determina y se vuelve a escribir en la memoria. Los parámetros específicos del modelo, como WEIGHT_BIAS y EPS, se pliegan en constantes.

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 con un núcleo de norma RMS regular, esta variante especializada ofrece un rendimiento y una latencia de 2 a 3 veces mejores tanto para el prellenado como para la decodificación.

El envío y la combinación de MoE reorganizan (shuffle) los tokens de entrada hacia los expertos a los que están enrutados, realizando una acumulación ponderada después de que los núcleos GEMM de MoE realizan las multiplicaciones de matrices relevantes. El rendimiento es de particular interés, ya que varios productos dependen de modelos MoE de varios tamaños. Según los tipos de despliegue, distinguimos diferentes núcleos para optimizar:

TP=1: Núcleo local para una sola GPU. El envío reorganiza los tokens hacia los expertos relevantes y la combinación los acumula. Se utiliza principalmente para despliegues de Qwen3 30B en H100.

DP=1, TP ≠ 1: Núcleo de un solo nodo y paralelismo de tensores. La entrada se replica, el envío selecciona los tokens para el rango actual y la combinación los acumula en todos los rangos. Acumulación a través de NVLink o NVLS. Utilizado para prellenadores grandes.

DP ≠ 1, NVLink: Núcleo con paralelismo de datos y paralelismo de tensores opcional. Cada rango envía tokens a cada uno de los otros rangos que alojan expertos específicos. La información de enrutamiento se agrega y los tokens se transfieren a través de NVIDIA NVLink. Utilizado para despliegues de H200 o GB200 con paralelismo de datos.

DP ≠ 1, InfiniBand: Núcleo peer-to-peer con paralelismo de datos y paralelismo de tensores que opera a través de NVLink y NVIDIA Quantum InfiniBand. Utilizado para despliegues de Hopper multinodo.

El núcleo local cuenta el número de tokens asignados a cada experto en función de la tabla de enrutamiento y asigna un índice único a cada token dentro de la lista de tokens pertenecientes a cada experto. Después del paso de enrutamiento, los tokens se copian en búferes opcionalmente rellenados (padded) para pasarlos a un núcleo GEMM de MoE. Utilizando las ubicaciones de los tokens, el núcleo de combinación realiza una acumulación en registros, escribiendo los tokens finales en un búfer de salida.

La fase de envío varía entre el prellenado y la decodificación: para pocos tokens, hay poca información de enrutamiento para leer, por lo que la procesamos en un solo bloque y señalizamos el resto de la cuadrícula, con bloques individuales copiando fragmentos de tokens. La asignación de tokens se determina incrementando atómicamente contadores en la memoria compartida. Dado que el prellenado debe procesar más información de enrutamiento, el enrutamiento se divide en múltiples bloques en un núcleo separado. Los índices se asignan primero localmente utilizando contadores en la memoria compartida y luego se agregan en todos los bloques en la memoria global para determinar la posición de cada token. La cuadrícula de este núcleo se limita al número de SM para poder utilizar barreras a nivel de cuadrícula. Un núcleo posterior realiza las copias, lanzado a través de un número mucho mayor de bloques y warps para utilizar plenamente el ancho de banda de memoria disponible.

El núcleo de combinación es relativamente simple, ya que realiza una acumulación vectorizada, basándose en la información de enrutamiento y en los desplazamientos calculados por el envío. CuTeDSL demostró ser valioso aquí, ya que el cálculo se pudo expresar a través de tensores CuTe simples que se compilan en código eficiente para todos los tipos de datos y dimensiones ocultas.

En un despliegue con paralelismo de tensores, la información de enrutamiento y los tokens de entrada se replican. El núcleo de envío local se puede reutilizar por completo con un ligero cambio para enrutar y seleccionar solo los tokens asignados a los expertos ubicados en el rango actual. Con algunos parámetros y especialización en tiempo de compilación, se reutilizan los núcleos de envío local.

Dado que el número de rangos suele ser menor que el número de expertos a los que se enruta un token, en el caso promedio se puede suponer que cada token tiene al menos una de sus réplicas enrutada a cada rango. En esta situación, el núcleo de combinación local se puede reutilizar para agregar las copias locales de un token, seguido de un all-reduce para sumar en todos los rangos participantes. En el caso del núcleo de prellenado, lanzamos núcleos separados; sin embargo, la decodificación fusiona ambas operaciones en un solo lanzamiento.

Para all-reduce, usamos NVLink plano si solo participan dos rangos, cambiando a NVLS cuando participan 4 o más rangos para un aumento del 10 al 20% en el rendimiento. Gracias al DSL, podemos compartir la mayor parte de la lógica de reducción y los intercambios de búferes de pares entre el MoE y los núcleos all-reduce regulares. Esta es una situación en la que el control cercano sobre el hardware es beneficioso, ya que se requieren operaciones DSL sobre multimem.ld_reduce y multimem.st para NVLS. La implementación de la reducción en sí sigue en gran medida el ejemplo all_reduce de las pruebas NCCL.

Los núcleos de paralelismo de datos se basan en NVLink para agregar tokens para múltiples dispositivos. En la fase de envío, antes de intercambiar tokens, los pares primero enrutan localmente para identificar el desplazamiento de destino en cada rango para cada token que están enviando. Posteriormente, comienzan a copiar los tokens en los búferes de recepción asignados para cada remitente para evitar conflictos. Las escrituras de liberación-almacenamiento (release-store) por token se utilizan para señalar la finalización. En el extremo receptor, la información de enrutamiento se agrega para poder condensar todos los tokens en un solo búfer para el GEMM de MoE. Los receptores sondean la bandera escrita por el par con una carga de adquisición (acquire-load) para esperar a que se completen las escrituras.

En el lado de la combinación, los tokens se copian a través de NVLink a los rangos de origen, con una bandera por token que indica la recepción. El núcleo de combinación es similar a las variantes local y de mononodo, con ligeras diferencias en el direccionamiento y cierta lógica de sincronización adicional.

Construimos una biblioteca de contenedores (wrappers) de PTX para exponer cargas y almacenamientos de liberación-adquisición con ámbito tanto a nivel de sistema como de dispositivo para implementar este núcleo. Al mejorar la granularidad de las banderas de sincronización y asignar más bloques para enviar y recibir tokens a través de NVLink, logramos mejoras de 2 veces en comparación con nuestros núcleos CUDA anteriores.

Nuestros núcleos peer-to-peer lograron previamente un rendimiento de última generación en GPU Hopper emparejadas con adaptadores InfiniBand ConnectX-7. Desde entonces, hemos migrado los núcleos de combinación a CuTeDSL para obtener más ganancias de rendimiento. Dado que las versiones de CUDA ya dependían en gran medida de la especialización en tiempo de compilación, demostraron ser los cuellos de botella en el tiempo de compilación en nuestro repositorio, lo que nos obligó a omitir algunos parámetros. Al migrar al DSL compilado por JIT, nos especializamos en más parámetros y ajustamos las asignaciones de cuadrícula, introduciendo más diferencias entre el prellenado y la decodificación para obtener un aumento de aproximadamente el 3% en la latencia.

Desde la perspectiva de los núcleos ROSE, la compatibilidad con Blackwell se divide en dos categorías: núcleos independientes de la arquitectura y específicos de la arquitectura. Los núcleos específicos de la arquitectura, que dependen de los núcleos tensoriales, requieren un trabajo explícito para adaptarse. Sin embargo, gracias al uso del DSL, los núcleos independientes de la arquitectura se adaptan trivialmente.

Una GPU Blackwell tiene hasta 160 SM, en contraste con los 132 de Hopper. Para la mayoría de los núcleos, esto es una constante en tiempo de compilación y la lógica en el código del lanzador del lado del host lo toma en cuenta para determinar el núcleo óptimo (prellenado frente a decodificación) a ejecutar. Además, Blackwell también introduce instrucciones de carga de 256 bits, que se pueden emitir fácilmente aumentando el ancho del vector (última dimensión de los tensores CuTe) para mejorar la utilización del ancho de banda de memoria.

Blackwell introduce soporte de hardware para un número significativo de esquemas de cuantización (MXFP8, MXFP4, NVFP4). El soporte para estos se implementó principalmente ajustando los núcleos existentes para el escalado por bloques de FP8. El soporte JIT es crucial aquí, ya que se pueden admitir nuevos modos sin hacer explotar el espacio de configuración en tiempo de compilación. En el pasado, encontramos que Triton era incómodo para expresar reducciones horizontales en pocos elementos, sobre la mitad o un cuarto de un warp. Un control más detallado sobre las reducciones era crucial para lograr un rendimiento óptimo al calcular un valor máximo absoluto en 16 o 32 elementos para NVFP4 y MXFP8/MXFP4, respectivamente.

Como lenguaje relativamente nuevo, los LLMs tuvieron poca exposición al código CuTeDSL durante su entrenamiento previo. Además, el estilo de nuestro código actual, con núcleos optimizados para activaciones y normas, era ligeramente diferente de los ejemplos existentes de CUTLASS o Quack, que estaban fuertemente optimizados para TMA e instrucciones de multiplicación de matrices específicas de la arquitectura. Llegar a un punto en el que los LLMs pudieran ayudar en la escritura de núcleos fue un proceso gradual.

Construimos manualmente un corpus inicial de núcleos y funciones de utilidad, siguiendo patrones consistentes, antes de poder depender de agentes para la conversión de núcleos. A medida que el corpus progresaba, pudimos realizar traducciones aproximadas de CUDA a CuTeDSL, con un retrabajo manual significativo posterior. En particular, descubrimos que los modelos tenían dificultades con las conversiones de puntero a tensor y la vectorización, quedándose atascados intentando corregir errores de tipo al intentar compilar una prueba. Han mostrado poca comprensión de la naturaleza basada en cuadrículas de los núcleos, confundiendo la sincronización de hilos y cuadrículas. Lo más frustrante fue la tentación de alucinar traducciones inexistentes de bibliotecas específicas de CUDA en lugar de compilar frente a nuestros propios reemplazos de CuTeDSL. Con el tiempo, a medida que el corpus de núcleos CuTeDSL crecía, descubrimos que los LLMs pueden construir la mayor parte del código base (skeleton code) y rellenar los espacios en blanco, lo que nos permite centrarnos en el diseño de alto nivel.

En comparación con CUDA, encontramos que CuTeDSL ofrece un rendimiento en tiempo de compilación mucho mejor, lo que nos permite especializarnos en más parámetros y lograr mejoras de rendimiento. Aunque los núcleos son ligeramente más detallados (verbose) debido a la falta de ciertas abstracciones en comparación con Triton, encontramos que la complejidad está justificada, ya que un control más detallado sobre el hardware nos permite construir núcleos más rápidos.

En el momento de escribir este artículo, la brecha de rendimiento entre CUTLASS C++ y los núcleos CuTeDSL es mínima, lo que abre oportunidades emocionantes para componer y fusionar núcleos GEMM de alto rendimiento con comunicación entre dispositivos, normalización o activaciones. Al comprometernos con este DSL, podremos extraer un mayor rendimiento de nuestros dispositivos para mejorar el rendimiento (throughput) en nuestras cargas de trabajo, desde incrustaciones hasta 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}
}

Inferencia en Perplexity AI

¿Por qué CuTeDSL?

Experiencia de desarrollo

Gran espacio de configuración

Especialización en relleno previo y decodificación

Sincronización de cuadrícula

Núcleos CuTeDSL

Norma QK

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

Local

Mononodo (Singlenode)

Paralelismo de datos (Data Parallel)

Peer-to-Peer

Compatibilidad con Blackwell

Generación de núcleos asistida por IA

Conclusión y trabajo futuro

Referencias

Cita