CuTeDSL en Perplexity

En Perplexity, utilizamos el motor de servicio optimizado para el tiempo de ejecución (ROSE, por sus siglas en inglés) desarrollado internamente para servir modelos que van desde incrustaciones hasta LLM de un billón de parámetros. Situado detrás de todas las API, incluidas Sonar, Search y Embeddings, nuestro tiempo de ejecución fue creado para capacitar a los investigadores a explorar los modelos y arquitecturas más adecuados, mientras se despliegan rápidamente los mejores candidatos en producción. Para lograr esto, necesitamos ser capaces de implementar, desplegar y optimizar la inferencia para arquitecturas de vanguardia a un ritmo acelerado. En esta entrada 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 las GPU de NVIDIA.

Los modelos personalizados, creados por varios equipos en 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 construyó inicialmente para servir mejor a los modelos Llama personalizados tanto para la decodificación de lenguaje como para la clasificación, con una interfaz compatible con el servidor de inferencia NVIDIA Triton. Con el tiempo, se convirtió en un motor de inferencia completo que sirve a importantes LLM, así como a modelos de clasificación, puntuación, ranking y otros basados en transformadores para una amplia gama de equipos con el fin de impulsar los productos de Perplexity. El objetivo de diseño de ROSE es simple: con el tiempo, se invierte esfuerzo para alcanzar el máximo rendimiento, pero la prioridad principal es adaptarse a las arquitecturas y casos de uso en constante evolución de nuestros modelos.

El espacio de configuración del kernel entre tipos de datos, dimensiones de cabeza, constantes de tiempo de compilación y ajustes de lanzamiento.

A grandes rasgos, ROSE consiste en 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 el procesamiento por lotes de las 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 manejados por ellos. Para los LLM, el motor es responsable del procesamiento por lotes, la fragmentación 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 procesamiento por lotes en línea. La ejecución se optimiza de tal manera 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 tubería, incluidos los núcleos de incrustación, norma, MoE y activación, son parte de la biblioteca de núcleos de ROSE. Para soportar mejor el hardware en el que ejecutamos la inferencia, ROSE está equipado con sus propias primitivas de comunicación entre dispositivos para el enrutamiento, despacho y combinación de MoE. Si bien los otros núcleos son relativamente simples, deben estar minuciosamente especializados para lograr el máximo rendimiento 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 variables de varios modelos (sesgo de peso frente a sin sesgo de peso para la norma RMS, funciones de activación, etc.). En consecuencia, decidimos utilizar 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 reside en expresar qué cálculo se debe realizar, sino en cómo debe hacerlo exactamente el hardware. Para construir los mejores núcleos, requerimos un lenguaje que ofrezca las abstracciones correctas 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, mientras 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 dificultan la velocidad de desarrollo. De los lenguajes que hemos encuestado o con los que hemos trabajado (incluidos CUDA, cuTile, TVM IR y Triton), consideramos que es el que mejor se adapta al desarrollo de nuestros núcleos.

En general, aunque no es tan de alto nivel y sucinto 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 máximo rendimiento en todos los casos extremos. Con CUDA, las oportunidades para la 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. Por el contrario, con un DSL compilado justo a tiempo, podemos especializarnos de forma más agresiva sin sacrificar el ritmo de iteración. A diferencia de TVM IR, que usaba CUDA como paso intermedio, descubrimos que la tubería basada en MLIR que bajaba a PTX era más robusta y más fácil de depurar que una que emitía CUDA como paso intermedio. Además, con el uso de operaciones personalizadas para envolver las 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 suele dedicarse a 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 explotar a varios minutos para un núcleo no trivial. Este cuello de botella se acentúa por el hecho de que las oportunidades para la compilación paralela surgen solo más tarde, a nivel de IR y PTX, con el manejo de plantillas de C++ ocurriendo secuencialmente. Si bien existen trucos para solucionar esto, al mover las instanciaciones de plantillas a unidades de compilación separadas, requieren ajustes significativos a los sistemas de compilación y el uso de macros difíciles de manejar. Por el contrario, los núcleos CuTeDSL que se compilan justo a tiempo (JIT) no requieren reglas de compilación complejas y no dificultan 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 más adelante pueden cubrir todo el espacio de diseño del núcleo.

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

El alto costo de compilación de CUDA se acentúa por la dificultad de rastrear el origen de los errores cuando se utiliza el despacho y la especialización basados en plantillas, como es el caso de CuTe y CUTLASS. Con un DSL, los errores pueden señalarse a la línea de código que los generó, facilitando significativamente la depuración. Además, es mucho más simple asignar el código generado a la instrucción DSL que lo originó a través de etapas intermedias de MLIR o confiando 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 los modelos, aún deben estar equipados con una gama de ajustes de configuración que pueden tener un impacto en el rendimiento. Los parámetros relevantes incluyen:

Dimensión oculta/Dimensión de cabezal/Número de cabezales de atención

Tipos de datos de entrada/salida

Esquema de cuantización

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

La arquitectura ROSE, incluidos su motor, modelo, biblioteca de kernels 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 siempre tomada o nunca tomada. Además, compilar un núcleo para un conjunto específico de ajustes de configuración resulta en menos instrucciones PTX/SASS, que son más fáciles de analizar y depurar cuando se usa cuda-gdb durante una investigación.

Para un tamaño de lote y una configuración determinados, se puede estimar el tráfico esperado a 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 del 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:

Prefill: tamaño de lote pequeño (pocos fragmentos, secuencias), muchos tokens por secuencia. El rendimiento debe maximizarse tanto para las incrustaciones como para los tokens de entrada de los LLM. Latencia en decenas a cientos de milisegundos.

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 objetivo de borrador, predicción de múltiples tokens). Latencia cercana a la latencia de lanzamiento del núcleo.

Para GEMM y atención, la etapa de prefill se ejecuta en una gran cantidad de tokens y el rendimiento está limitado por el rendimiento de cálculo 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 pesos de DRAM en los multiprocesadores de transmisión y núcleos tensoriales, limitados 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 prefill y decode 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 prefill y decode si las diferencias algorítmicas o las asignaciones de cuadrícula/bloque beneficiarían a cualquiera de ellos.

El parámetro que es universalmente diferente entre los núcleos de prefill y decode 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 de 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 un 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, 64K en Hopper y Blackwell. Si se conoce el número de hilos con antelación, el compilador sabe exactamente cuántos registros puede asignar a cada hilo y puede lograr compilar 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 tiene que ser recompilado 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 prefill y decode. Para utilizar completamente el ancho de banda de memoria de una GPU, no es suficiente asignar siempre un gran número de warps por bloque, ya que el ancho de banda completo solo puede lograrse 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 vectorial por hilo ideal (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 requerido por el núcleo. Primero intentamos maximizar el número de SMs: si hay menos warps que SMs, se elige un solo warp por bloque. Esta suele ser la variante decode del núcleo. Si hay suficientes warps, aumentamos gradualmente el número de warps a 8, 16 o 32, comparando 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. Especialmente cuando se trata de reducciones horizontales de sub-warp, descubrimos que CuTeDSL nos ofrecía el control detallado necesario para generar código óptimo.

En algunas situaciones, para distribuir el trabajo a través de tantos bloques como sea posible, se puede usar cooperative_groups::this_grid().sync() para sincronizar todos los bloques participantes. Se aplican restricciones: dado que un bloque no puede ser desprogramado 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 SMs disponibles en el dispositivo, para asegurar que en cualquier momento todos los bloques puedan actualizar y sondear las ubicaciones de memoria global involucradas en la sincronización. Esta restricción se aplica mediante cudaLaunchCooperativeKernel. El costo de una barrera de cuadrícula es de aproximadamente 3us 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 decode, evitar la sobrecarga vale la pena.

Si una de las cargas de trabajo requiere más bloques que el número de SMs, el núcleo debe dividirse, basándose en la barrera implícita de los lanzamientos de núcleo. Este es principalmente el caso para el prefill. 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 decode que operan con 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 se ilustra en el ejemplo a continuación. Con una estrategia híbrida de este tipo, los núcleos decode pueden optimizarse para alcanzar latencias mínimas y los núcleos prefill pueden optimizarse para un rendimiento máximo. A diferencia de las soluciones basadas en CUDA, esta lógica puede definirse de forma más sucinta 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 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 auto-atención de los recientes modelos Qwen, se aplica por cabezal, 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 entre warps, con la posibilidad de ajustar múltiples cabezales en un solo warp. Este se convirtió en un caso de uso donde pudimos optimizar agresivamente en tiempo de compilación para encontrar las mejores asignaciones de token-a-warp para todas las dimensiones de cabezal y tipos de datos.

La norma RMS es una operación bastante directa, 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 los cabezales de atención suelen ser pequeños 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 bfloat16. En consecuencia, un medio warp es suficiente para procesar un cabezal, y la norma al cuadrado se puede determinar mediante una reducción horizontal realizada por shuffle_sync_bfly. Partiendo de la dimensión del cabezal y el tipo de datos, podemos determinar el número de cabezales por warp y especificar los argumentos y el número de rondas de reducción en el momento de la compilación.

Podemos entonces construir un núcleo especializado con relativamente pocas líneas de código, confiando en cargas/almacenamientos vectorizados automáticamente y utilizando operaciones de arquitectura específica de nivel inferior para la reducción de warp. El núcleo determina en qué cabezal (Q o K) trabajar e identifica los pesos correspondientes a aplicar. Carga un vector completo del cabezal y primero suma localmente. Posteriormente, según el número de hilos por cabezal, 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, están plegados a 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 2-3 veces mejores tanto para el prefill como para el decode.

El despacho y la combinación de MoE barajan los tokens de entrada a los expertos a los que están enrutados, realizando una acumulación ponderada después de que los núcleos MoE GEMM realicen las multiplicaciones de matriz relevantes. El rendimiento es de particular interés, ya que varios productos dependen de modelos MoE de diversos tamaños. Según los tipos de despliegue, distinguimos diferentes núcleos para optimizar:

TP=1: Núcleo local para una sola GPU. El despacho baraja los tokens a los expertos relevantes y la combinación los acumula. Utilizado principalmente para despliegues Qwen3 30B en H100.

DP=1, TP ≠ 1: Kernel paralelo a tensorial, de un solo nodo. La entrada se replica, el despacho 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 prefillers grandes.

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

DP ≠ 1, InfiniBand: Núcleo de paralelismo de datos y paralelismo tensorial punto a punto que opera a través de NVLink y NVIDIA Quantum InfiniBand. Utilizado para despliegues Hopper de múltiples nodos.

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 que pertenecen a cada experto. Después del paso de enrutamiento, los tokens se copian en búferes opcionalmente acolchados para ser pasados a un núcleo MoE GEMM. Utilizando las ubicaciones de los tokens, el núcleo de combinación realiza una acumulación en los registros, escribiendo los tokens finales en un búfer de salida.

La fase de despacho varía entre prefill y decode: para pocos tokens, hay poca información de enrutamiento para leer, por lo que lo procesamos en un solo bloque y señalamos 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 prefill debe procesar más información de enrutamiento, el enrutamiento se divide en varios bloques en un núcleo separado. Los índices se asignan primero localmente usando 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 SMs para poder usar 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 completamente el ancho de banda de memoria disponible.

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

En un despliegue paralelo a tensorial, la información de enrutamiento y los tokens de entrada se replican. El núcleo de despacho local se puede reutilizar completamente 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, los núcleos de despacho local se reutilizan.

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 asumir 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 prefill, lanzamos núcleos separados, sin embargo, el decode fusiona ambas operaciones en un solo lanzamiento.

Para all-reduce, usamos NVLink simple si solo participan dos rangos, cambiando a NVLS cuando participan 4 o más rangos para un aumento del 10-20% en el rendimiento. Gracias al DSL, podemos compartir la mayor parte de la lógica de reducción y los intercambios de búfer entre pares entre los núcleos MoE y los núcleos all-reduce regulares. Esta es una situación donde 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í misma sigue en gran medida el ejemplo de all_reduce de las pruebas de NCCL.

Los núcleos de paralelismo de datos dependen de NVLink para agregar tokens para múltiples dispositivos. En la fase de despacho, antes de intercambiar tokens, los pares primero enrutan localmente para identificar el desplazamiento objetivo en cada rango para cada token que envían. Posteriormente, comienzan a copiar los tokens en búferes de recepción asignados para cada remitente para evitar conflictos. Se utilizan escrituras de liberación-almacenamiento por token para señalar la finalización. En el lado del receptor, la información de enrutamiento se agrega para poder condensar todos los tokens en un solo búfer para MoE GEMM. Los receptores sondean la bandera escrita por el par con una carga-adquisición 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 un solo nodo, con ligeras diferencias en el direccionamiento y alguna lógica de sincronización adicional.

Construimos una biblioteca de envoltorios PTX para exponer cargas-almacenamientos de liberación-adquisición con alcance 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 2x sobre nuestros núcleos CUDA anteriores.

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

Desde la perspectiva de los núcleos ROSE, el soporte para Blackwell cae en dos categorías: núcleos independientes de la arquitectura y núcleos 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 de tiempo de compilación y la lógica en el código de lanzamiento del lado del host lo tiene en cuenta para determinar el núcleo óptimo (prefill vs decode) a ejecutar. Además, Blackwell también introduce instrucciones de carga de 256 bits, que pueden emitirse 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 bloque FP8. El soporte JIT es crucial aquí, ya que los nuevos modos se pueden admitir sin explotar el espacio de configuración de 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 fue 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 LLM tuvieron poca exposición al código CuTeDSL durante su preentrenamiento. 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 y multiplicaciones de matrices específicas de la arquitectura. Llegar a un punto en el que los LLM 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 que pudiéramos confiar en agentes para la conversión de núcleos. A medida que avanzaba el corpus, podíamos 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 atrapados tratando de corregir errores de tipo cuando intentaban que una prueba se compilara. Han mostrado poco entendimiento de la naturaleza basada en cuadrícula de los núcleos, confundiendo la sincronización de hilos y cuadrícula. Lo más frustrante es que hubo una tentación de alucinar traducciones inexistentes de bibliotecas específicas de CUDA en lugar de construir sobre nuestros propios reemplazos de CuTeDSL. Con el tiempo, a medida que crecía el corpus de núcleos CuTeDSL, descubrimos que los LLM pueden construir la mayor parte del código esqueleto y completar 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 mucho mejor en tiempo de compilación, lo que nos permite especializarnos en más parámetros y lograr mejoras de rendimiento. Si bien los núcleos son ligeramente más detallados debido a la falta de ciertas abstracciones en comparación con Triton, encontramos que la complejidad está justificada ya que el control 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 los núcleos CUTLASS C++ y CuTeDSL es mínima, lo que abre interesantes oportunidades 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 en nuestras cargas de trabajo, desde incrustaciones hasta 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 qué CuTeDSL?

Experiencia del desarrollador

Gran espacio de configuración

Especialización en Relleno previo y descodificación

Sincronización de cuadrícula

Núcleos CuTeDSL

Norma QK

Despacho/Combinación de Mixture-of-Experts

Local

Un solo nodo

Paralelo a datos

Punto a punto

Soporte para Blackwell

Generación de núcleos asistida por IA

Conclusión y trabajo futuro

Referencias

Citación