El agente con preguntas superó al agente con herramientas: 113 % frente a 50 % de la velocidad de FlashAttention-3

Los kernels de GPU rápidos para attention (atención) son la base de la velocidad con la que se entrenan y ejecutan los grandes modelos, pero escribirlos a mano lleva mucho tiempo: cada generación de tarjetas gráficas exige reescribirlos de nuevo, y a los equipos expertos les cuesta meses. Investigadores de Carnegie Mellon University comprobaron si se puede encargar este trabajo a un agente de programación y llegaron a una conclusión poco intuitiva. Lo que ayuda a un agente atascado no son nuevas herramientas ni más libertad, sino una pregunta corta en el momento oportuno.

En una misma tarea, el agente con el conjunto completo de herramientas se detuvo en el 50% de la velocidad de FlashAttention-3 (el kernel de atención de referencia para las GPU Hopper) tras 133 pasos. El agente al que se le hicieron preguntas llegó al 113% en las mismas 26 configuraciones. Para un desarrollador, esto significa lo siguiente: si el agente se atasca, lo más probable es que le falte una hipótesis sobre dónde buscar el problema, y no otra porción de datos.

Por qué las herramientas no bastan

Los agentes actuales para optimizar kernels se sitúan en los dos extremos de una escala. Los agentes basados en guion fijan con rigidez cada paso: una política entrenada, un menú de optimizaciones, roles fijos. Mientras la receta encaja, son económicos; cuando no encaja, se atascan.

Los agentes libres mutan por sí mismos un conjunto de kernels, pero pagan esa libertad con miles de ejecuciones, ediciones repetidas y candidatos erróneos. En un trabajo anterior, un agente de este tipo exploró en siete días más de 500 direcciones y publicó 40 versiones, y necesitó un supervisor para salir de los bucles de ediciones fallidas. Alrededor del 30% de las líneas que añaden agentes de este tipo coinciden byte a byte con otras eliminadas anteriormente.

Ante un objetivo difícil, ambas escuelas responden igual: dar más al agente. Una salida del profiler ya procesada, un análisis del desensamblador, búsqueda en la documentación, un almacén de «lecciones aprendidas». Los autores consideran que el agente ya tiene suficiente información. Los modelos se aferran con avidez a la primera solución, se detienen de forma prematura, dan vueltas entre variantes casi idénticas e informan de que han terminado cuando la verificación demuestra lo contrario.

Cómo cambia una pregunta la tarea

La orden «hazlo más rápido» es abierta, y el modelo responde a ella con avidez, retocando lo que tiene más a mano. La pregunta «¿estás ganando en ciclos del kernel o solo en tiempo de reloj?» está acotada, y para responderla el agente desmontará y reescribirá una parte del kernel que por sí mismo no habría tocado. Una buena pregunta no contiene la respuesta ya hecha: nombra una hipótesis que el agente no había considerado, y se puede comprobar con medios ordinarios, midiendo.

Del código específico, el método solo incluye la infraestructura de verificación. Un solo comando compila al candidato, compara sus salidas con las del kernel original y con una referencia en fp32, mide la velocidad en cada escenario respecto al kernel original y calcula la aceleración como media geométrica. El objetivo se define como una condición de parada verificable, no como una petición al modelo.

Las preguntas llegan cuando el progreso se detiene, cuando el agente declara «este es el límite» o «está listo», y tras una ganancia sospechosamente grande. Al principio las formula una persona, y las preguntas que se repiten se acumulan en un banco que el agente utiliza por sí mismo. El método también se integró en una búsqueda evolutiva, donde las indicaciones para las mutaciones se formulan como preguntas.

Qué resultados se obtuvieron

Las campañas completas produjeron kernels de paso hacia delante de atención para Hopper y Blackwell: un único archivo CUDA sin CUTLASS, CuTe, cuBLAS ni cuDNN. Son correctos en los 182 escenarios de cada arquitectura (en Blackwell, en todos los ejecutables). En media geométrica son 1,16 veces más rápidos que FlashAttention-3, es decir, un 16% más, y 1,46 veces más rápidos que FlashAttention-4, un 46% más. En las formas de consulta de producción, la ventaja sobre FlashAttention-4 fue de 1,10 veces. Los autores los consideran los kernels abiertos más rápidos de su clase en los escenarios comprobados.

Las estadísticas de las intervenciones resultan curiosas: 162 mensajes humanos y 11,5 mil llamadas a herramientas en 16 sesiones. Las preguntas no hicieron que el agente usara menos las herramientas; simplemente las utilizó de otra manera. Cada gran avance siguió a una pregunta que llevó al agente a otra capa del sistema o que rechazó el límite que él mismo había declarado, y ninguno siguió a un mensaje que se limitara a aportar información.

Las salvedades son evidentes. Los autores construyeron ellos mismos el agente autónomo de comparación, siguiendo el modelo de trabajos anteriores. El experimento es uno solo: en él existe una referencia sólida con la que medir. Parte del trabajo de orientación lo realiza una persona. Los investigadores publicaron los kernels, la infraestructura de verificación, el banco de preguntas y todas las trazas.

Conclusión práctica: antes de conectar una herramienta más o una capa de búsqueda adicional, defina un objetivo verificable en todos los escenarios y prepare preguntas para los momentos de estancamiento, de límite declarado y de victoria demasiado grande.