CUDA Agent generación de kernels CUDA es el objetivo de un sistema de aprendizaje por refuerzo desarrollado por ByteDance Seed y Tsinghua AIR. La propuesta utiliza un modelo de lenguaje dentro de un entorno real de desarrollo para escribir, comprobar y perfilar kernels destinados a GPU, con el foco puesto en superar el rendimiento de los kernels generados por compiladores.
El sistema parte de Seed1.6, un modelo propietario de tipo MoE con 23.000 millones de parámetros activos y 230.000 millones en total. En KernelBench, Seed1.6 alcanzaba una tasa de resolución del 74 %, pero solo lograba superar a torch.compile en el 27,2 % de las tareas. CUDA Agent eleva ese resultado hasta el 96,8 % de kernels más rápidos que el compilador y un rendimiento medio geométrico de 2,11 veces sobre él.
CUDA Agent obtuvo un 98,8 % de tareas superadas y un 96,8 % de resultados más rápidos que torch.compile en las 250 pruebas de KernelBench, con una media geométrica de 2,11 veces sobre el compilador.
Un agente que programa, mide y corrige kernels
La diferencia frente a una generación de código convencional está en el ciclo de trabajo. CUDA Agent opera con herramientas de terminal y edición de archivos, ejecuta compilaciones en un entorno aislado y comprueba tanto la corrección como el rendimiento. El flujo sigue un patrón ReAct y utiliza instrucciones distribuidas mediante un archivo SKILL.md.
El sistema perfila primero el modelo de PyTorch y después modifica model_new.py para incorporar kernels personalizados. La iteración continúa hasta intentar superar en al menos un 5 % a torch.compile, con tolerancias de 1e-2 para el error absoluto y relativo. El entorno bloquea permisos y restringe determinadas vías de atajo para reducir resultados artificiales.
Controles contra la optimización engañosa
El entrenamiento incorpora verificaciones sobre cinco entradas aleatorias, sincronización del dispositivo durante el perfilado y calentamiento previo. También impide recurrir a funciones alternativas de torch.nn.functional mediante gestores de contexto y no ofrece una herramienta de búsqueda web.
La recompensa no se basa directamente en una proporción de velocidad. Asigna valores discretos: -1 si falla la corrección, 3 cuando el kernel supera en más de un 5 % tanto al modo eager como a torch.compile, 2 si solo mejora el modo eager y 1 en el resto de casos. En las ablaciones, sustituir esta recompensa por una proporción bruta de velocidad reduce la tasa de mejora frente al compilador del 96,8 % al 60,4 %.
Los mejores resultados aparecen al fusionar operadores
El equipo generó el conjunto CUDA-Agent-Ops-6K a partir de operadores de las bibliotecas torch y transformers. El proceso combina hasta cinco clases de operadores en una capa fusionada y descarta ejemplos no deterministas, constantes, demasiado similares a KernelBench o fuera del rango de ejecución establecido. El conjunto final contiene 6.000 muestras, de las que el 83,77 % son composiciones de dos operadores.
El nivel centrado en secuencias de operadores fue el más sólido: logró un 100 % de tareas resueltas y una tasa del 100 % de resultados más rápidos que torch.compile, con una media geométrica de 2,80 veces. En el nivel más complejo, la tasa de resolución fue del 94 %, la de mejora frente al compilador del 90 % y la media geométrica de 1,52 veces.
Entre los casos estudiados figuran una multiplicación matricial diagonal reescrita como escalado por filas, con una mejora de 73,31 veces; una cadena de multiplicación, división, suma y escalado fusionada y reordenada, con 24,04 veces; y un bloque básico de ResNet con BatchNorm integrado en la convolución, con 3,59 veces.
Los pesos siguen sin estar disponibles
El resultado no equivale a un lanzamiento abierto del agente entrenado. Los pesos de CUDA Agent no se han publicado, y su entrenamiento utilizó 128 GPU NVIDIA H20 solo para el entorno de perfilado. Esto sitúa una replicación completa más cerca de laboratorios con grandes recursos, proveedores de GPU en la nube y equipos de infraestructura de gran tamaño.
Sí están disponibles el conjunto CUDA-Agent-Ops-6K, la especificación SKILL.md y las recetas de calentamiento y recompensa. Estos componentes podrían servir a equipos medianos que quieran adaptar el enfoque a un modelo base abierto, aunque el material no confirma un modelo concreto ni un paquete de reproducción completo.
Qué puede aportar a la infraestructura de IA
La técnica está orientada a cargas donde los kernels fusionados afectan directamente a la latencia o al coste: servidores de inferencia, nubes de GPU, conducción autónoma, trading cuantitativo, imagen médica y sistemas de recomendación. También puede ser útil para volver a ajustar kernels cuando cambia la generación de GPU o para reducir el coste por token en servicios de inferencia.
Las cifras proceden de KernelBench y de evaluaciones del propio trabajo, por lo que todavía queda pendiente comprobar su comportamiento en cargas de producción y con una reproducción independiente. Además, la ausencia de los pesos limita el acceso directo al agente que obtuvo estos resultados.
Conclusión
CUDA Agent muestra que la generación de código CUDA puede beneficiarse de un ciclo que combine escritura, ejecución, verificación y perfilado, especialmente en problemas de fusión de operadores. Sus resultados en KernelBench son elevados, pero el sistema sigue siendo parcialmente cerrado: la comunidad puede estudiar el conjunto de datos y las recetas, no ejecutar el modelo entrenado tal como se evaluó.
Encuentra más IA en nuestra sección de IA o suscríbete a Guslok – Guías 100% en YouTube para más contenido.
Sobre el Autor
