KernelArc: A Multi-Agent Framework for GPU Kernel Optimization
KernelArc es un marco de trabajo multi-agente que optimiza autónomamente kernels de GPU mediante agentes paralelos y especializados en estrategias que se coordinan a través de memoria compartida y guardas deterministas, logrando los primeros puestos en la clasificación de SOL-ExecBench para diversos flujos de trabajo en GPUs NVIDIA H100 y B200.
Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo
Las computadoras modernas dependen de chips especializados llamados unidades de procesamiento gráfico, o GPUs, para manejar los cálculos masivos requeridos por la inteligencia artificial. Estos chips son increíblemente potentes, pero también son máquinas complejas con muchas partes diferentes que deben trabajar en perfecta sincronía. Para sacar el máximo provecho de ellos, los ingenieros deben escribir programas diminutos y altamente específicos llamados kernels que le dicen al chip exactamente cómo mover datos y realizar operaciones matemáticas. Durante años, este ha sido un trabajo para expertos humanos que pasan incontables horas ajustando el código, equilibrando el uso de la memoria y cronometrando las operaciones para exprimir cada bit de velocidad. A medida que estos chips se han vuelto más sofisticados, el trabajo manual se ha vuelto más difícil y la brecha entre lo que el hardware puede hacer y lo que el software logra se ha ampliado.
Los investigadores han comenzado a utilizar modelos de lenguaje de gran tamaño —programas informáticos entrenados en vastas cantidades de texto— para ayudar a escribir y mejorar este código. Estos modelos pueden sugerir cambios, probarlos y aprender de los resultados, actuando como un ingeniero automatizado. Sin embargo, un único agente automatizado suele estancarse. Puede encontrar una buena solución y seguir puliéndola, perdiendo otras aproximaciones mejores que se encuentran en una dirección diferente. Es como un excursionista que encuentra un buen sendero para subir una montaña y sigue caminando por él, sin darse cuenta de que un sendero diferente cercano conduce a una cima mucho más alta. Para resolver esto, un equipo de investigadores en IMEC ha desarrollado un nuevo sistema llamado KernelArc, que utiliza un grupo de estos agentes automatizados trabajando juntos para explorar muchos caminos a la vez.
El sistema KernelArc opera asignando diferentes estrategias a varios agentes que se ejecutan en paralelo. En lugar de que un solo agente intente resolver un problema por su cuenta, cada agente del grupo se enfoca en un ángulo de ataque diferente, como cambiar la forma en que se almacenan los datos, alterar la precisión matemática utilizada o fusionar múltiples pasos en uno solo. Estos agentes no comparten todo su proceso de pensamiento ni cada paso que dan, lo que sería demasiado desordenado y confuso. En su lugar, comparten solo sus conclusiones finales: qué funcionó y qué falló. Escriben estos resultados en un espacio de memoria compartida que actúa como un tablero de anuncios. Si un agente descubre un truco que acelera un cálculo, publica el resultado. Otros agentes pueden leer esa publicación y usar la idea para guar de su propio trabajo, evitando callejones sin salida y construyendo sobre el éxito.
Para asegurar que los agentes no pierdan el tiempo con código roto, el sistema incluye un árbitro automatizado estricto. Este árbitro somete cada nueva sugerencia de código a una serie de pruebas para verificar si es correcto y qué tan rápido es. Si una sugerencia falla la prueba, se descarta inmediatamente. Si funciona pero no es más rápida que la mejor versión actual, se mantiene como una copia de seguridad pero no reemplaza al líder. Solo cuando una nueva versión es tanto correcta como más rápida, se convierte en el nuevo estándar. Este proceso permite al equipo de agentes explorar una amplia gama de posibilidades sin perderse en los detalles de cada intento. El sistema también tiene un mecanismo de seguridad que se activa si un agente se queda estancado en una meseta, donde no se pueden encontrar más mejoras. Cuando esto sucede, el sistema obliga al agente a intentar un enfoque completamente diferente, asegurando que la búsqueda continúe avanzando.
Los investigadores probaron este sistema en dos de las GPUs más avanzadas disponibles, la NVIDIA H100 y la B200. Se enfocaron en un conjunto de tareas estándar utilizadas para medir el rendimiento, que incluyen operaciones como la multiplicación de matrices, mecanismos de atención utilizados en modelos de lenguaje y diversas formas de fusión de datos. En una prueba específica que involucraba a un solo agente trabajando solo con una guía detallada, el sistema logró alcanzar una velocidad de 766 teraflops, lo que es aproximadamente un 3.2 por ciento más rápido que la mejor biblioteca existente para esa tarea específica. Esto demostró que un solo agente podía profundizar mucho en un problema específico si se le daba un camino claro. Sin embargo, cuando los investigadores cambiaron al sistema multiagente KernelArc para abordar un conjunto más amplio de tareas, los resultados fueron aún más impresionantes. El sistema produjo implementaciones personalizadas para diversas operaciones complejas, incluyendo mecanismos de atención especializados y capas fusionadas para modelos de lenguaje de gran tamaño.
En un tablero público que clasifica la velocidad de estas operaciones a través de muchas formas y tamaños de datos, las entregas de KernelArc ocuparon el primer lugar en varias categorías. Para una tarea que involucra atención y adición residual, un solo agente se estancó en una puntuación de rendimiento de 0.441. El sistema multiagente, al compartir conocimientos y explorar diferentes direcciones, rompió esa barrera y alcanzó una puntuación de 0.481. En otra prueba que involucraba una tarea de atención compleja, el sistema multiagente logró una aceleración de casi 291 veces en comparación con una implementación de referencia estándar, y más de 143 veces más rápido que una base altamente optimizada. Estos resultados sugieren que tener múltiples agentes compartiendo sus hallazgos permite al sistema encontrar mejores soluciones más rápido que un solo agente trabajando solo, especialmente cuando el espacio del problema es grande y variado.
El estudio también analizó cómo se comporta el sistema cuando cambian el número de agentes y la cantidad de memoria compartida. Encontraron que dar a los agentes una memoria compartida donde pudieran publicar sus éxitos y fracasos ayudó a alcanzar resultados más sólidos dentro de un número fijo de intentos. Cuando la memoria era ilimitada, el sistema funcionó mejor, alcanzando una aceleración de más de 290 veces en comparación con el punto de partida. Esto indica que la capacidad de aprender de la experiencia colectiva del grupo es un factor clave para el éxito. Los investigadores señalaron que el valor de cada característica, como la memoria compartida o la especialización de estrategias, depende de la tarea específica y de la etapa de la búsqueda. A veces los agentes necesitan explorar ampliamente, y otras veces necesitan refinar una solución específica.
Este trabajo no pretende haber resuelto todos los problemas de optimización informática, ni sugiere que los ingenieros humanos ya no sean necesarios. Los resultados son específicos para las tareas y el hardware probados, y el sistema todavía depende de un marco diseñado por humanos para guiar a los agentes. Sin embargo, los hallazgos demuestran que un grupo coordinado de agentes automatizados puede explorar un rango más amplio de soluciones que un solo agente. Al compartir solo sus conclusiones más valiosas, estos agentes pueden evitar repetir errores y construir sobre los descubrimientos de los demás. Este enfoque ofrece una forma prometedora de manejar la creciente complejidad de la computación moderna, donde la brecha entre el potencial del hardware y el rendimiento del software continúa creciendo. El sistema demuestra que, con la coordinación adecuada, las herramientas automatizadas pueden ayudar a desbloquear todo el poder de los chips que impulsan la próxima generación de la inteligencia artificial.
¿Ahogado en artículos de tu campo?
Recibe resúmenes diarios de los artículos más novedosos que coincidan con tus palabras clave de investigación — con resúmenes técnicos, en tu idioma.