En los últimos años, los modelos de aprendizaje automático se han vuelto una parte esencial del software cotidiano, y gran parte de su tiempo de ejecución depende de un puñado de núcleos computacionales críticos, como la multiplicación de matrices, la convolución o la normalización. Optimizar estos kernels es una de las formas más directas de reducir la latencia y los costos operativos, pero hacerlo tradicionalmente ha requerido de ingenieros especializados que escriban código GPU de bajo nivel a mano.

Ahora, los sistemas basados en agentes y modelos de lenguaje grandes (LLM) prometen automatizar esta tarea, aunque las herramientas actuales adolecen de varias limitaciones: se evalúan sobre tensores aleatorios y kernels aislados, producen código CUDA que debe ser reintegrado manualmente por los desarrolladores, se centran principalmente en modelos LLM de PyTorch y ofrecen poco soporte para inspección o depuración.

El nuevo proyecto Kernel Forge, presentado en arXiv, busca resolver estas limitaciones con un enfoque integral y de código abierto. Se trata de un sistema autónomo de extremo a extremo que acepta cualquier modelo PyTorch sin necesidad de modificaciones, compatible con cargas de trabajo de visión, difusión y modelos de lenguaje. Kernel Forge utiliza un algoritmo de búsqueda de árbol de Monte Carlo (MCTS) para explorar múltiples rutas de optimización en lugar de seguir una única cadena de refinamiento lineal, lo que aumenta la variedad de soluciones exploradas. Además, incluye una interfaz gráfica que permite a los usuarios monitorear el progreso, examinar kernels candidatos y depurar fallos de forma intuitiva.

La evaluación del equipo se realizó sobre cuatro modelos PyTorch representativos—un modelo de visión (ResNet-50), un modelo de difusión (Stable Diffusion 3.5 Medium), un modelo de lenguaje pequeño (Gemma 4 E2B) y un modelo de lenguaje grande (Qwen 3.5 35B-A3B)—ejecutados en una NVIDIA DGX Spark equipada con GPU GB10. Con solo 50 iteraciones de optimización por kernel, Kernel Forge logró superar el modo eager de PyTorch en 14 kernels diferentes. Los resultados más destacados incluyen un aumento de rendimiento de 1.52× en adaptive_avgpool2d de ResNet-50, 1.70× en group_norm de Stable Diffusion 3.5 Medium, 2.83× en softmax de Gemma 4 E2B y 1.54× en softmax de Qwen 3.5 35B-A3B.

Estos avances tienen implicaciones importantes tanto para la investigación como para la industria. Por un lado, democratizan el acceso a la optimización de alto rendimiento en GPU, permitiendo que equipos sin una experiencia profunda en programación de bajo nivel puedan obtener ganancias de velocidad significativas. Por otro lado, el enfoque basado en MCTS abre nuevas posibilidades para la optimización multiobjetivo y la exploración de diseños más creativa, más allá de la mera iteración lineal. La interfaz gráfica, por su parte, aborda una necesidad crítica de visibilidad y control, algo que muchas herramientas de optimización anteriores dejaban en segundo plano.

Aunque los resultados son prometedores, también plantean interrogantes sobre la generalización del método a otras arquitecturas de GPU, la escalabilidad a modelos aún más grandes y la robustez frente a diferentes patrones de entrada. Futuras iteraciones de Kernel Forge podrían integrar módulos de verificación estadística más avanzados y soportar automáticamente conversiones a otras APIs de computación acelerada, como AMD ROCm o Intel oneAPI.

En resumen, Kernel Forge representa un paso significativo hacia la automatización real de la optimización de kernels en GPU. Al combinar un agente basado en LLM con una estrategia de búsqueda rica y una interfaz de usuario intuitiva, la herramienta no solo acelera el desarrollo de aplicaciones de IA de alto rendimiento, sino que también reduce la barrera de entrada para la participación en el ecosistema de la computación acelerada. Su código abierto garantiza que la comunidad pueda contribuir a su evolución, lo que potencialmente acelerará aún más el ritmo de innovación en este espacio crítico.