ByteDance Seed y Tsinghua AIR presentan CUDA Agent: un sistema RL agente a gran escala para la generación de kernel CUDA

ByteDance Seed y Tsinghua AIR han lanzado CUDA Agent, un sistema de aprendizaje por refuerzo agente que entrena un modelo de lenguaje grande para escribir núcleos de GPU que superen a un compilador. La brecha a la que apunta es estrecha pero obstinada: los modelos de frontera ya producen CUDA correcto, solo producen CUDA lento. En KernelBench, el modelo base Seed1.6 supera el 74,0 % de las tareas, pero supera a torch.compile en sólo el 27,2 % de ellas, con una aceleración media geométrica de 0,69 veces, lo que significa que sus núcleos son, en promedio, más lentos que lo que el compilador genera por sí solo. CUDA Agent cierra esa brecha colocando el modelo dentro de un entorno de desarrollo CUDA real con creación de perfiles, comprobaciones de corrección y una zona de pruebas con permisos bloqueados, y luego entrenándolo con PPO para 150 pasos en un contexto de 131,072 tokens. El resultado es una tasa de aprobación del 98,8% y una tasa de compilación 96,8% más rápida que la antorcha en el punto de referencia de 250 tareas, con una media geométrica de 2,11 veces sobre la compilación, aproximadamente 40 puntos por delante de Claude Opus 4.5 y Gemini 3 Pro en la división de nivel 3 más difícil.

¿Es desplegable?

En parte, pero el agente capacitado no es liberado. Está construido sobre Seed1.6, un modelo patentado de MoE con 23B activos y 230B de parámetros totales, y el documento no incluye pesos. Público: el conjunto de datos CUDA-Agent-Ops-6K, la especificación SKILL.md y las recetas de recompensa y calentamiento.

Qué empresas: Solo la zona de pruebas de creación de perfiles utilizó 128 GPU NVIDIA H20, lo que permite la replicación completa dentro de laboratorios de vanguardia, nubes de GPU y grandes equipos de infraestructura. Los equipos de tamaño mediano aún pueden adoptar las partes (conjunto de datos, recompensa por hitos, restricciones contra la piratería de recompensas, especificaciones de habilidades) además de un modelo base abierto.

Industrias y aplicaciones: infraestructura de inteligencia artificial y servicio de inferencia, nube de GPU, conducción autónoma, comercio cuantitativo, imágenes médicas y sistemas de recomendación, en cualquier lugar donde los núcleos fusionados se encuentren en un camino de latencia crítica. Los usos incluyen fusionar secuencias de operadores que torch.compile maneja mal, reducir el costo por token y reajustar los núcleos en todas las generaciones de GPU.

Síntesis de datos

El equipo de investigación rastrea operadores de referencia de las bibliotecas de antorchas y transformadores. Luego, un LLM toma muestras de hasta cinco clases de operadores de antorcha y las apila en una capa fusionada. Un filtro mantiene solo los operadores que se ejecutan tanto en modo ansioso como de compilación, son deterministas, producen resultados no constantes y se ejecutan entre 1 ms y 100 ms en modo ansioso. Se eliminan las muestras con similitud AST superior a 0,9 con cualquier tarea de KernelBench. El resultado es CUDA-Agent-Ops-6K: 6.000 muestras, el 83,77% de ellas composiciones de dos operadores.

Medio ambiente y recompensa

El bucle del agente refleja las herramientas de OpenHands (Bash, Read/Write, Edit/MultiEdit, Glob, Grep, NotebookEdit, BashOutput, KillBash) bajo un patrón ReAct. Las instrucciones CUDA se envían en el formato Agent Skills. SKILL.md le dice al modelo que cree un perfil del modelo PyTorch, reescriba model_new.py con kernels personalizados, compile en un sandbox de GPU e itere hasta que el kernel sea al menos un 5 % más rápido que torch.compile en atol=1e-2, rtol=1e-2.

El pirateo de recompensas tiene cinco contramedidas: scripts de verificación y creación de perfiles con permisos bloqueados, administradores de contexto que prohíben las alternativas funcionales de torch.nn, verificaciones con cinco entradas aleatorias, creación de perfiles con sincronización y calentamiento de dispositivos, y ninguna herramienta de búsqueda web.

La recompensa es discreta en lugar de una relación de aceleración bruta. r ∈ {−1, 1, 2, 3}: −1 en caso de error de corrección, 3 si el núcleo borra tanto ansioso como torch.compile en más del 5%, 2 si borra solo ansioso, 1 en caso contrario.

https://arxiv.org/pdf/2602.24286v1

Resultados

Tabla 1, en general: tasa de aprobación del 98,8 %, 98,4 % más rápido que ansioso, 96,8 % más rápido que torch.compile, a 2,60 × y 2,11 × geomedia respectivamente. El nivel 2 (secuencias de operador) es la división más fuerte: 100 % de aprobación, velocidad 100 % más rápida, 2,80 veces más que torch.compile. El nivel 3 alcanza un 94,0 % de aprobación, una velocidad un 90,0 % más rápida y 1,52 ×, aproximadamente 40 puntos por encima de Claude Opus 4.5 (50,0 %) y Gemini 3 Pro (52,0 %) en una velocidad más rápida en comparación con la compilación.

Una inconsistencia: el resumen y la introducción indican tasas 100%/100%/92% más rápidas para los niveles 1 a 3, mientras que la Tabla 1 informa 97,0%/100,0%/90,0%. La tabla 1 es la tabla principal de resultados.

Las ablaciones son contundentes. La eliminación del bucle del agente reduce la tasa más rápida en comparación con la compilación del 96,8 % al 14,1 %. Una recompensa bruta de aceleración da un 60,4%, ningún RFT da un 49,8% más colapso de recompensa, un preentrenamiento sin valor da un 50,9% más trayectorias desbocadas.

Los estudios de caso muestran lo que aprende la política. Un matmul diagonal reescrito como escalado por filas: 73,31 × sobre torch.compile. Una cadena matmul-división-suma-escala reordenada y fusionada: 24,04×. Un ResNet BasicBlock con BatchNorm plegado en convolución ycudnnConvolutionBiasActivationForward: 3,59 ×.