Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
Kernel Forge es un harness agéntico de extremo a extremo y de código abierto que aprovecha los Modelos de Lenguaje de Gran Escala y la Búsqueda en Árbol de Monte Carlo para generar y optimizar automáticamente kernels de CUDA para diversas cargas de trabajo de PyTorch, logrando aceleraciones significativas sobre el modo eager de PyTorch mientras proporciona una interfaz gráfica para la inspección y depuración.
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
Imagina que estás intentando ejecutar un videojuego en tu computadora. A veces, el juego se ejecuta sin problemas, pero otras veces se traba o tiene lag. Esto sucede porque la computadora tiene que realizar millones de cálculos matemáticos diminutos cada segundo para dibujar las imágenes y manejar la física. En el mundo de la inteligencia artificial (IA), estos cálculos son aún más intensos. Para que una IA pueda "pensar", depende de instrucciones especiales de alta velocidad llamadas kernels. Piensa en un kernel como una receta específica y súper eficiente que un chef utiliza para picar vegetales. Si el chef usa un cuchillo desafilado y un método lento, la cena se retrasará. Si usa un cuchillo afilado y especializado, la comida estará lista al instante.
Durante años, escribir estas "recetas" para la IA ha sido como intentar escribir una novela en un idioma que apenas hablas. Requiere de expertos altamente capacitados para codificar a mano instrucciones complejas para los chips de la computadora (GPUs) para hacer que funcionen más rápido. Pero ahora, tenemos un nuevo tipo de ayudante: los Modelos de Lenguaje Extensos (LLM). Tal vez conozcas a los chatbots inteligentes que pueden escribir historias o resolver problemas matemáticos. Los científicos están enseñando ahora a estos chatbots a escribir las "recetas" (kernels) para la IA, con la esperanza de que puedan hacerlo más rápido y mejor que los humanos. Sin embargo, hay un inconveniente: el hecho de que un chatbot escriba una receta que se vea bien en el papel no significa que realmente cocinará la comida más rápido en una cocina real.
Aquí es donde entra un nuevo proyecto llamado Kernel Forge. Investigadores de la Universidad de Michigan construyeron un sistema inteligente que no solo le pide a un chatbot que escriba código; actúa como un entrenador de tiempo completo, un inspector de calidad y un gerente de proyectos, todo en uno. Querían ver si la IA podía realmente tomar el control de la optimización de estas recetas para modelos de IA del mundo real, no solo para casos de prueba ficticios y creados artificialmente.
El Problema: La Trampa del "Aislamiento"
Imagina que estás entrenando a un robot para correr. Si solo lo pruebas en una pista plana y vacía, podría parecer un velocista de clase mundial. Pero en el momento en que lo pones en una calle de la ciudad con baches y gente, podría tropezar y caerse. Esto es exactamente lo que sucedió con los intentos anteriores de usar IA para escribir kernels de GPU. La mayoría de las herramientas probaban la IA dándole problemas matemáticos aleatorios y ficticios en una pista vacía. Generaban una pieza de código, la probaban de forma aislada y decían: "¡Mira qué rápido es esto!".
Pero en el mundo real, los modelos de IA son como ciudades bulliciosas. Las "recetas" (kernels) tienen que trabajar con formas específicas de datos, cantidades específicas de memoria y tienen que convivir armoniosamente con las otras recetas que están justo al lado. Una receta que es rápida en una prueba aleatoria podría ser lenta o incluso fallar cuando se usa dentro de un modelo de IA real, como los que generan imágenes o charlan contigo. Las herramientas anteriores solían dejar en manos de los desarrolladores humanos el descubrir cómo encajar estas nuevas recetas escritas por IA dentro de la compleja máquina, lo cual era tedioso y propenso a errores.
La Solución: Kernel Forge
Los investigadores crearon Kernel Forge, una herramienta de código abierto que actúa como un puente entre el escritor de IA y el mundo real. En lugar de pedirle a la IA que adivine, Kernel Forge primero observa un modelo de IA real (como uno que reconoce rostros o escribe historias) ejecutándose en una computadora. Toma notas sobre exactamente qué "recetas" se están utilizando, qué tan grandes son los datos y cuánto tiempo tardan en ejecutarse.
Una vez que tiene estos datos del mundo real, le entrega el trabajo a un agente de IA (un bot inteligente impulsado por un modelo de lenguaje extenso). Pero esto no es un proceso de un solo paso. El sistema utiliza una estrategia ingeniosa llamada Búsqueda de Árbol Monte Carlo (Monte Carlo Tree Search). Imagina a un detective resolviendo un misterio. En lugar de seguir solo una corazonada, el detective intenta muchos caminos diferentes. Si un camino lleva a un callejón sin salida, no se rinde; regresa y prueba un ángulo diferente. Del mismo modo, Kernel Forge no solo escribe una versión de una receta y espera lo mejor. Genera muchas versiones, las prueba y, si una es lenta, intenta arreglarla o intenta un enfoque completamente diferente. Mantiene un "árbol genealógico" de todos los intentos, recordando cuáles fueron prometedores incluso si no fueron perfectos al principio.
Los Resultados: Victorias Rápidas, Pero No en Todas Partes
El equipo probó Kernel Forge en cuatro tipos diferentes de modelos de IA: uno para reconocer imágenes (ResNet-50), uno para generar arte (Stable Diffusion 3.5 Medium) y dos para charlar y razonar (Gemma 4 y Qwen 3.5). Ejecutaron estas pruebas en una computadora potente con una GPU NVIDIA GB10.
Esto es lo que encontraron:
- Funciona, pero es selectivo: El sistema generó con éxito nuevas recetas más rápidas para muchas partes de la IA. En algunos casos, el código escrito por la IA fue significamente más rápido que el código estándar que la computadora suele utilizar. Por ejemplo, en el generador de imágenes, el nuevo código para la "normalización de grupo" (group normalization) fue 1.70 veces más rápido. En el chatbot Gemma 4, el nuevo código para "softmax" (un paso matemático utilizado para tomar decisiones) fue un masivo 2.83 veces más rápido.
- La red de seguridad "protegida": El sistema es muy cuidadoso. Tiene un "guardián" que revisa cada nueva receta. Si la nueva receta escrita por la IA es más lenta o comete un error, el sistema cambia inmediatamente de vuelta al código original y confiable. Nunca impone una mala receta al usuario. Esto significa que, incluso si la IA intenta optimizar una parte del código y falla, la computadora no se bloquea ni se ralentiza; simplemente utiliza el método antiguo y seguro.
- Los "Grandes Jugadores" son difíciles de vencer: Los investigadores notaron algo interesante. Las partes de la IA que ocupan la mayor parte del tiempo (los "grandes jugadores") suelen estar ya gestionadas por código muy maduro y escrito por expertos de empresas como NVIDIA. La IA tuvo dificultades para superar a estos expertos. Por ejemplo, en el generador de imágenes, la operación "lineal" (que ocupa más del 50% del tiempo) fue en realidad más lenta cuando la IA intentó reescribirla. El sistema decidió sabiamente mantener el código original para esa parte.
- Las pequeñas victorias se acumulan: Las mayores aceleraciones ocurrieron en las partes más pequeñas y menos críticas del código. Aunque la IA no pudo vencer a los expertos en las tareas más grandes, encontró formas ingeniosas de acelerar las tareas más pequeñas entre 1.5 y 2.8 veces.
El Costo de la Curiosidad
Los investigadores también analizaron cuánto costó este "pensamiento". Utilizaron un modelo de IA potente para generar el código, y cada vez que la IA escribía una línea de código o revisaba un resultado, costaba una pequeña cantidad de dinero (basado en el uso de la API). Descubrieron que a medida que permitían que la IA probara más y más variaciones (hasta 50 rondas de intento y corrección), el costo aumentaba. Sin embargo, la velocidad adicional que obtenían no siempre coincidía con el dinero extra gastado. Esto sugiere que, si bien la IA puede encontrar atajos excelentes, debemos ser inteligentes sobre cuándo detener la búsqueda, especialmente si la parte del código que intentamos arreglar no es muy importante para la velocidad general.
La Conclusión
Kernel Forge demuestra que estamos entrando en una nueva era donde la IA puede ayudar a escribir el código de bajo nivel que hace que otras IA funcionen más rápido. No es una varita mágica que lo resuelve todo instantáneamente; aún no puede vencer fácilmente a los mejores expertos humanos en las tareas más grandes. Pero es una herramienta poderosa que puede encontrar mejoras de velocidad ocultas en los detalles más pequeños, todo mientras mantiene el sistema seguro y estable.
Los investigadores lanzaron su herramienta como código abierto, lo que significa que cualquiera puede usarla para intentar hacer que sus propios modelos de IA sean más rápidos. Demostraron que, al combinar la creatividad de la IA con la seguridad de las pruebas del mundo real, podemos comenzar a optimizar los motores de nuestro mundo digital sin necesidad de tener un doctorado en ciencias de la computación para hacerlo. Es un paso hacia un futuro donde nuestras computadoras no solo sean más inteligentes, sino también mucho más eficientes.
¿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.