Introducción
El equipo Seed de ByteDance, en colaboración con el Instituto de Investigación en Inteligencia Artificial (AIR) de la Universidad de Tsinghua, ha lanzado CUDA Agent, un sistema de aprendizaje por refuerzo con agentes a gran escala diseñado para enseñar a los modelos de lenguaje a generar kernels CUDA de alto rendimiento.
Este trabajo aborda una debilidad de larga data en la generación de código de bajo nivel por parte de la IA. Los modelos de lenguaje de vanguardia suelen generar código CUDA que compila y se ejecuta correctamente, pero para cargas de trabajo GPU de nivel de producción, la corrección por sí sola no es suficiente. Los kernels generados también deben competir con las implementaciones generadas por compiladores altamente optimizados.
CUDA Agent se construye precisamente en torno a este segundo problema: no solo generar código CUDA válido, sino aprender a perfilar, verificar y optimizar iterativamente los kernels hasta lograr mejoras significativas en el tiempo de ejecución.
Por qué generar código CUDA correcto no es suficiente
Antes de introducir el aprendizaje por refuerzo con agentes, el modelo base Seed1.6 ya demostraba una capacidad considerable para escribir código CUDA en KernelBench.
De los 250 problemas de KernelBench Nivel 1–3 utilizados por los investigadores, Seed1.6 logró una tasa de aprobación del 74.0%. En otras palabras, podía generar soluciones funcionalmente correctas para la mayoría de las tareas de referencia del modelo base.
Pero el rendimiento era otra cuestión.
Solo el 27.2% de los kernels generados por el modelo base eran más rápidos que torch.compile, con una velocidad media geométrica de solo 0.69× en comparación con la línea base del compilador.
| Modelo | Tasa de aprobación | Más rápido que torch.compile | Velocidad media geométrica frente a torch.compile |
|---|---|---|---|
| Modelo base Seed1.6 | 74.0% | 27.2% | 0.69× |
| CUDA Agent | 98.8% | 96.8% | 2.11× |
Esta brecha pone de manifiesto el principal desafío en la generación automática de kernels GPU.
Los modelos de lenguaje pueden comprender la sintaxis de CUDA, pero aún pueden producir accesos a memoria ineficientes, lanzamientos excesivos de kernels, malas elecciones de particionado en bloques, e intermedios innecesarios.
Perfilado de rendimiento
- Guía de optimización específica para CUDA
- Entorno de sandbox restringido
- Señales de recompensa vinculadas a los resultados reales de ejecución
El archivo SKILL.md también restringe los atajos que podrían distorsionar los resultados de las pruebas de referencia. Por ejemplo, el agente no puede simplemente recurrir a operaciones arbitrarias de PyTorch dentro de sus implementaciones personalizadas de kernels.
Esto es importante porque los sistemas de aprendizaje por refuerzo pueden encontrar formas de maximizar la recompensa sin resolver el problema de optimización previsto.
El entrenamiento con PPO extiende al agente a trayectorias de optimización largas
Los investigadores utilizaron optimización de política proximal (PPO) para entrenar a CUDA Agent.
Un desafío principal es que la optimización de GPU es inherentemente una tarea de horizonte temporal largo. El modelo puede necesitar pasar por múltiples rondas de edición de código, compilación, depuración, perfilado y optimización adicional antes de alcanzar el resultado deseado.
Por lo tanto, el artículo adopta diferentes longitudes de contexto y fases de entrenamiento para estabilizar el aprendizaje.
El modelo base es Seed1.6, un modelo de mezcla de expertos con 230 mil millones de parámetros totales y 23 mil millones de parámetros activos.
Para el aprendizaje por refuerzo con agentes:
- La ventana de contexto es de 131,072 tokens.
- El búfer de repetición de entrenamiento permite hasta 150 rondas de interacción del agente.
- La evaluación permite hasta 200 rondas de interacción del agente.
- El modelo se entrena con 150 pasos de RL.
- Los modelos actor y crítico emplean una estrategia de precalentamiento en múltiples fases antes de la optimización completa de horizonte largo.
Esto es más complejo que simplemente ajustar un modelo sobre un par de indicaciones y respuestas de CUDA.
El objetivo es enseñar al modelo cómo mejorar sus propios kernels mediante la interacción repetida con la retroalimentación de ejecución.
El sandbox separa la compilación del perfilado de GPU
La medición fiable del rendimiento es fundamental, porque la velocidad de ejecución forma parte de la recompensa.
Por ello, los investigadores construyeron una arquitectura de sandbox con recursos CPU–GPU desacoplados.
El sandbox de terminal basado en Docker maneja tareas orientadas a CPU (como la compilación), mientras que la verificación y el perfilado de rendimiento se envían a un grupo de sandboxes GPU dedicado que contiene 128 GPU NVIDIA H20.
Esta separación tiene un doble propósito.
Primero, aísla la compilación y la interacción del agente de las pruebas comparativas en GPU. Segundo, la asignación dedicada de GPU reduce la interferencia entre cargas de trabajo concurrentes, haciendo que las mediciones de latencia sean más estables.
Esto es importante para el aprendizaje por refuerzo: si las mediciones de tiempo están contaminadas por ruido, el modelo recibe señales de recompensa poco fiables y puede aprender comportamientos de optimización incorrectos.
CUDA Agent alcanza una tasa de aprobación del 98.8% y supera al compilador en el 96.8% de las tareas
El sistema final se evaluó en 250 tareas de KernelBench Niveles 1 a 3.
CUDA Agent logró:
- Tasa de aprobación general del 98.8%
- 98.4% más rápido que la ejecución inmediata de PyTorch
- 96.8% más rápido que
torch.compile - Aceleración media geométrica de 2.60× frente a la ejecución inmediata
- Aceleración media geométrica de 2.11× frente a
torch.compile
Estos resultados representan una mejora sustancial respecto al punto de partida de Seed1.6.
El modelo no solo mejoró en la generación de código que pasa las pruebas de corrección, sino que los kernels que genera también tienen más probabilidades de superar las líneas base del compilador.
Resultados por nivel de dificultad de KernelBench
Las mejoras de rendimiento se mantienen sólidas en los tres niveles de dificultad de KernelBench.
| Dificultad KernelBench | Tasa de aprobación | Proporción más rápida que torch.compile | Aceleración media geométrica frente a torch.compile |
|---|---|---|---|
| Nivel 1 | 100.0% | 97.0% | 1.87× |
| Nivel 2 | 100.0% | 100.0% | 2.80× |
| Nivel 3 | 94.0% | 90.0% | 1.52× |
El rendimiento en el Nivel 2 es particularmente destacable.
Estas tareas implican secuencias de operadores donde las oportunidades de optimización mediante fusión y movimiento de memoria no siempre son aprovechadas eficazmente por las heurísticas estáticas de los compiladores.
Los autores del artículo sostienen que el optimizador iterativo basado en aprendizaje puede explorar un espacio más amplio de estrategias de implementación, incluidos patrones de acceso a memoria específicos del hardware, particionado en bloques y elecciones de fusión.
Los datos de entrenamiento se publicarán próximamente
Actualmente, los pesos completos del modelo entrenado no están incluidos en la versión pública del proyecto.
Sin embargo, los investigadores han publicado varios componentes importantes del stack de entrenamiento.
Crea un sitio de presentacion y capta leads en minutos
Describe tu idea una vez y We0 AI puede generar un sitio de presentacion, paginas y CMS, y ayudarte a atraer clientes y trafico tras el lanzamiento.
Una generación completa de proyectos para registro gratuito
Lo mejor para probar un flujo de generación completo y ver rápidamente el primer borrador del proyecto.
CUDA-Agent-Ops-6K
El conjunto de datos CUDA-Agent-Ops-6K contiene 6,000 tareas sintéticas de entrenamiento a nivel de operadores.
Su proceso de construcción incluye:
- Recopilar operadores semilla de bibliotecas como
torch,transformers, etc. - Usar LLMs para combinar múltiples operadores en tareas de fusión más complejas.
- Filtrar las tareas generadas mediante verificaciones basadas en resultados de ejecución.
La fase de filtrado elimina casos que son aleatorios, demasiado simples, inválidos o demasiado similares a las tareas de evaluación.
El conjunto de datos está publicado en Hugging Face bajo la licencia CC BY 4.0.
SKILL.md y entorno del agente
El repositorio de GitHub también incluye el archivo de instrucciones SKILL.md orientado a CUDA y el entorno de trabajo del agente.
Estos materiales documentan el flujo de trabajo de optimización, las herramientas disponibles, las restricciones y la configuración del entorno de ejecución utilizada para guiar al agente.
Mecanismo de recompensa y esquema de entrenamiento
El artículo y el repositorio describen el diseño de recompensas, las fases de precalentamiento, la inicialización del crítico y el esquema de entrenamiento basado en PPO para estabilizar el proceso de optimización de largo horizonte.
Esto es particularmente útil para los investigadores centrados en el entrenamiento de agentes para programación de sistemas, más allá de simplemente reproducir las puntuaciones finales de referencia.
Por qué esto es importante más allá de KernelBench
La optimización de kernels CUDA sustenta una gran parte de la infraestructura moderna de IA.
Los kernels más rápidos pueden mejorar:
- El rendimiento del entrenamiento de modelos
- La latencia de inferencia de LLMs
- La utilización de GPU
- La eficiencia de costos del servicio
- Los pipelines de IA en tiempo real
- Las cargas de trabajo de computación numérica de alto rendimiento
El artículo original de AIBase señala su importancia en el campo de la IA.
Aplicaciones potenciales en áreas como infraestructura, servicios de inferencia de modelos de gran escala, conducción autónoma y trading cuantitativo — donde diferencias mínimas de latencia pueden ser cruciales.
Los resultados de CUDA Agent no implican que los compiladores tradicionales estén obsoletos.
torch.compile sigue siendo una línea base automática sólida, y la evaluación de CUDA Agent depende de entornos de referencia controlados y de un entorno de GPU específico.
Las conclusiones que muestra esta investigación, aunque de alcance más limitado, siguen siendo importantes: con retroalimentación de ejecución suficiente y un aprendizaje por refuerzo especializado, los modelos de lenguaje pueden aprender estrategias de optimización que superan la línea base de compiladores estáticos en la gran mayoría de las tareas de núcleos GPU evaluadas.
Esto significa transformar la ingeniería de rendimiento en un campo razonable para el aprendizaje agéntico, y no solo para la generación de código en un solo intento.
Preguntas frecuentes
¿Qué es CUDA Agent?
CUDA Agent es un sistema de aprendizaje por refuerzo a gran escala para agentes inteligentes, desarrollado por ByteDance Seed, el Instituto de Investigación de Industria Inteligente (AIR) de la Universidad de Tsinghua y su SIA-Lab conjunto. Entrena modelos de lenguaje para escribir, verificar, analizar y optimizar núcleos CUDA de forma iterativa.
¿En qué modelo se basa CUDA Agent?
La investigación utiliza Seed1.6 como modelo base. El artículo lo describe como un modelo de mezcla de expertos con un total de 230 mil millones de parámetros y 23 mil millones de parámetros activados.
¿Qué es KernelBench?
KernelBench es un punto de referencia abierto que evalúa si los modelos de lenguaje pueden generar núcleos GPU correctos y eficientes para programas de PyTorch. CUDA Agent fue evaluado en 250 tareas que abarcan los niveles 1 a 3 de KernelBench.
¿Cuánto más rápido es CUDA Agent que torch.compile?
En todo el conjunto de referencia, CUDA Agent logra una aceleración media geométrica de 2.11 veces en comparación con torch.compile. Sus núcleos generados son más rápidos que la línea base del compilador en el 96.8% de las tareas evaluadas.
¿CUDA Agent utiliza aprendizaje por refuerzo?
Sí. El sistema utiliza PPO, combinado con un calentamiento en múltiples etapas, preentrenamiento del modelo de valor, diseño robusto de recompensas y trayectorias largas de agentes en múltiples rondas.
¿Cuánto puede durar una trayectoria de optimización de CUDA Agent?
El rollout de entrenamiento permite hasta 150 rondas de agente, mientras que la evaluación permite hasta 200 rondas. La ventana de contexto de entrenamiento del agente es de 131,072 tokens.
¿CUDA Agent es de código abierto?
El proyecto ha publicado su repositorio de GitHub, el entorno de agentes, SKILL.md, la receta de entrenamiento y el conjunto de datos CUDA-Agent-Ops-6K. Los pesos completos del modelo entrenado no están incluidos en los materiales publicados actualmente.
¿Qué es CUDA-Agent-Ops-6K?
CUDA-Agent-Ops-6K es un conjunto de datos que contiene 6,000 tareas sintéticas de entrenamiento CUDA a nivel de operadores. Está diseñado específicamente para el aprendizaje por refuerzo a gran escala de agentes e incluye pasos de filtrado para garantizar que las tareas sean ejecutables, deterministas, no triviales y estén separadas del conjunto de evaluación de KernelBench.
Herramientas relacionadas
- CUDA Agent: Página oficial del proyecto del sistema de generación de núcleos CUDA de ByteDance Seed y el Instituto de Investigación de Industria Inteligente de la Universidad de Tsinghua.
- CUDA Agent GitHub: Repositorio oficial que contiene el entorno de agentes,
SKILL.mdy los materiales publicados del proyecto. - CUDA-Agent-Ops-6K: Conjunto de datos oficial de entrenamiento con 6,000 muestras publicado con el proyecto.
- KernelBench: Punto de referencia abierto para evaluar núcleos GPU generados por modelos de lenguaje grandes.
- PyTorch: Marco de aprendizaje profundo cuya ejecución en modo eager y
torch.compileson líneas base clave en la investigación. - NVIDIA CUDA: Documentación oficial de NVIDIA sobre programación CUDA y desarrollo de núcleos GPU.
Enlaces relacionados
- Página del proyecto CUDA Agent: Descripción general oficial, resultados de referencia, diagramas del proyecto, artículo, código y enlaces a los datos.
- Artículo de CUDA Agent (arXiv): Artículo de investigación completo, "Aprendizaje por refuerzo agéntico a gran escala para la generación de núcleos CUDA de alto rendimiento".
- Repositorio de GitHub de CUDA Agent: Repositorio de código fuente con el entorno publicado y los archivos del proyecto.
- Conjunto de datos CUDA-Agent-Ops-6K: Conjunto de datos oficial de Hugging Face con 6,000 tareas de operadores sintéticos.
- Repositorio de GitHub de KernelBench: Entorno de referencia y evaluación para la generación eficiente de núcleos GPU.
- Documentación de PyTorch para
torch.compile: Documentación oficial de PyTorch de la línea base de compilador utilizada en la investigación. - Guía de programación CUDA C++ de NVIDIA: Guía oficial de NVIDIA para programación y optimización CUDA.
Resumen
CUDA Agent aborda una debilidad específica en el código de sistemas generado por modelos de lenguaje grandes: el código CUDA generado no solo debe ser correcto, sino también genuinamente más rápido que las alternativas producidas por compiladores.
Partiendo de Seed1.6, los investigadores utilizaron un entorno de agentes específico para CUDA, retroalimentación de ejecución, diseño robusto de recompensas y entrenamiento PPO de horizonte largo para aumentar la tasa de éxito del 74.0% al 98.8%, y la proporción de kernels más rápidos que torch.compile del 27.2% al 96.8%.
El proyecto también publicó un conjunto de datos de entrenamiento con 6,000 tareas, SKILL.md, materiales del entorno de agentes y la receta de entrenamiento, proporcionando una base concreta para que los investigadores continúen trabajando en la dirección de optimización de GPU aprendida.
La principal contribución de CUDA Agent es demostrar que la ingeniería de rendimiento puede aprenderse mediante un bucle agéntico iterativo — y no solo aproximarse mediante la generación de código en un solo intento.
