El Portal de las Tecnologías para la Innovación

Automatización de la generación de kernel de GPU con DeepSeek-R1 y escala de tiempo de inferencia

A medida que los modelos de IA amplían sus capacidades para resolver desafíos más sofisticados, está surgiendo una nueva ley de escala conocida como escalamiento en tiempo de prueba o escalamiento en tiempo de inferencia.

También conocida como razonamiento de IA o pensamiento a largo plazo , esta técnica mejora el rendimiento del modelo al asignar recursos computacionales adicionales durante la inferencia para evaluar múltiples resultados posibles y luego seleccionar el mejor, red neuronal. Esto permite que la IA diseccione y resuelva sistemáticamente problemas complejos de una manera similar a cómo los humanos diseccionan problemas complejos y los resuelven individualmente para llegar a una solución final.

En esta publicación, hablamos sobre un experimento realizado por ingenieros de NVIDIA que utilizaron uno de los modelos de código abierto más nuevos, el modelo DeepSeek-R1, junto con potencia computacional adicional durante la inferencia para resolver un problema complejo. El experimento consistía en generar automáticamente núcleos de atención de GPU que fueran numéricamente correctos y optimizados para diferentes tipos de atención sin ninguna programación explícita. 

Los resultados resultaron ser mejores que los núcleos optimizados desarrollados por ingenieros expertos en algunos casos. 

La necesidad de núcleos de atención optimizados y los desafíos asociados

La atención es un concepto clave que revolucionó el desarrollo del modelo de lenguaje grande (LLM). Es un mecanismo poderoso que permite a los modelos de IA centrarse selectivamente en las partes más relevantes de la información de entrada al realizar tareas. Al centrarse en la información importante, la operación de atención ayuda a los modelos a realizar mejores predicciones y encontrar patrones ocultos en los datos. 

NVIDIA GTC2025

La complejidad computacional de la operación de atención crece cuadráticamente en relación con la longitud de la secuencia de entrada. Esto motiva la necesidad de desarrollar una implementación optimizada de nivel inferior (es decir, un núcleo de GPU) para evitar errores de ejecución que surgen de implementaciones simples (por ejemplo, errores de falta de memoria) y para fines de eficiencia computacional. 

Existen múltiples variantes de atención (causal, incrustaciones posicionales relativas, coartada, etc.) y a menudo los ingenieros deben utilizar una combinación de estas variantes para una tarea determinada.

Los modelos multimodales (por ejemplo, los transformadores de visión) introducen una capa adicional de desafíos, ya que requieren mecanismos de atención especializados (Atención al Vecindario Espacial) para mantener la información espacio-temporal que a menudo se encuentra en la visión por computadora, los modelos de generación de video, etc. 

Cuatro imágenes de personas observando saltos de esquí tienen diferentes valores de dilatación, de 1 a 4. La dilatación expande el cuadro delimitador de un solo esquiador (dilatación = 1) a casi toda la imagen (dilatación = 4).
Figura 1. Atención del vecindario en entradas 2D

Crear un kernel de GPU optimizado para la atención requiere mucha habilidad y tiempo, incluso para ingenieros de software experimentados.

Los últimos LLM como DeepSeek-R1 han demostrado ser muy prometedores en las tareas de generación de código, pero aún enfrentan desafíos para crear código optimizado en el primer intento. Esto hace que sea necesario usar otras estrategias en el momento de la inferencia para generar código optimizado. 

El siguiente mensaje es un ejemplo de entrada de usuario para un kernel de atención de incrustaciones posicionales relativas.

Please write a GPU attention kernel to support relative position encodings. Implement the relative positional encoding on the fly within the kernel. The complete code should be returned, including the necessary modifications.Use the following function to compute the relative positional encoding:def relative_positional(score, b, h, q_idx, kv_idx):    return score + (q_idx - kv_idx)When implementing the kernel, keep in mind that a constant scaling factor 1.44269504 should be applied to the relative positional encoding due to qk_scale = sm_scale * 1.44269504. The PyTorch reference does not need to scale the relative positional encoding, but in the GPU kernel, use:qk = qk * qk_scale + rel_pos * 1.44269504Please provide the complete updated kernel code that incorporates these changes, ensuring that the relative positional encoding is applied efficiently within the kernel operations.

En ocasiones, los LLM pueden producir código alucinado o mezclar sintaxis de distintos lenguajes o marcos, lo que provoca errores de código inmediatos o ineficiencias. Calcular la asignación óptima de subprocesos de GPU tampoco es una tarea trivial y desafiante, que a menudo requiere un refinamiento iterativo para lograr un núcleo correcto y eficiente.

Escalado del tiempo de inferencia para generar núcleos de GPU optimizados

Para obtener los mejores resultados con núcleos de atención optimizados, los ingenieros de NVIDIA crearon un nuevo flujo de trabajo que incluye un verificador especial junto con el modelo DeepSeek-R1 durante la inferencia en un modo de circuito cerrado durante una duración predeterminada. 

Un diagrama de flujo muestra el mensaje inicial que genera un núcleo de atención. DeepSeek-R1 crea código de GPU, que se verifica. Si no se cumplen los criterios, las GPU Hopper refinan y dan formato al mensaje. El resultado final son núcleos optimizados para GPU.
Figura 2. Escalado del tiempo de inferencia con DeepSeek-R1 en la plataforma NVIDIA Hopper

El flujo de trabajo se inicializa primero mediante un mensaje manual y el modelo DeepSeek-R1 genera el código de la GPU (es decir, el núcleo) en el primer paso. El verificador se ejecuta en una GPU NVIDIA H100. Analiza el núcleo generado y crea mensajes nuevos que se proporcionan como entrada al modelo DeepSeek-R1. 

Este enfoque de bucle cerrado mejora el proceso de generación de código al guiarlo de una manera diferente cada vez. El equipo descubrió que al dejar que este proceso continuara durante 15 minutos, se obtenía un núcleo de atención mejorado. 

Un gráfico de barras que muestra la aceleración promedio del núcleo de atención en la GPU Hopper, compara la aceleración de diferentes tipos de núcleo de atención entre dos enfoques: "PyTorch API (Flex Attention)" en naranja y "NVIDIA Workflow con DeepSeek-R1" en verde. La API de PyTorch mantiene una línea base de 1x para todos los núcleos, mientras que NVIDIA Workflow con DeepSeek-R1 logra aceleraciones de 1,1x para Causal Mask y Document Mask, 1,5x para Relative Positional, 1,6x para Alibi Bias y Full Mask, y 2,1x para Softcap.
Figura 3. Rendimiento de los núcleos de atención optimizados generados automáticamente con atención flexible

Este flujo de trabajo produjo núcleos numéricamente correctos para el 100 % de los problemas de nivel 1 y el 96 % de los problemas de nivel 2, según lo probado por el benchmark KernelBench de Stanford .

La tasa de resolución de nivel 1 en KernelBench se refiere a la métrica numérica correcta que se utiliza para evaluar la capacidad de los LLM de generar núcleos de GPU eficientes para tareas computacionales específicas. Esta prueba es parte de una serie de desafíos para probar las capacidades de los últimos LLM en programación de GPU.

La figura 4 muestra cómo el presupuesto de tiempo de inferencia afecta la tasa de resolución del agente. Asignar más de 10 minutos por problema en la categoría de Nivel 1 permite que el flujo de trabajo produzca código numéricamente correcto para la mayoría de los 100 problemas.

Un gráfico de líneas muestra la cantidad de núcleos numéricamente correctos generados a lo largo del tiempo. La línea se acerca al 95 % a los 10 minutos y al 100 % a los 20 minutos.
Figura 4. Resultados del escalamiento del tiempo de inferencia en núcleos de GPU optimizados

Núcleos de GPU optimizados en DeepSeek-R1

Estos resultados muestran cómo se puede utilizar el último modelo DeepSeek-R1 para generar mejores núcleos de GPU utilizando más potencia de procesamiento durante el tiempo de inferencia. Esta sigue siendo una nueva área de investigación con resultados preliminares sobre un enfoque prometedor que genera automáticamente núcleos de atención efectivos. 

Si bien hemos tenido un buen comienzo, es necesario seguir trabajando para generar mejores resultados de manera consistente para una mayor variedad de problemas. Estamos entusiasmados con los recientes avances en DeepSeek-R1 y su potencial. 

Para obtener más información o comenzar, consulte el microservicio NIM DeepSeek-R1 , ahora disponible en build.nvidia.com .

Recursos relacionados

NVIDIA Blog. T. Ch., B. X. y K. D. Traducido al español

Artículos relacionados

Scroll al inicio