La eficiencia computacional en sistemas distribuidos a escala de hyperscaler depende críticamente de la optimización de kernels de bajo nivel que traducen operaciones de modelos de IA a instrucciones específicas de hardware. Sin embargo, la proliferación de arquitecturas de hardware (NVIDIA, AMD, MTIA, CPUs), la evolución constante de modelos de Machine Learning (desde DLRMs hasta LLMs) y la necesidad de operadores personalizados más allá de las bibliotecas estándar, han creado un problema combinatorio de optimización de kernels que excede la capacidad de ajuste manual por expertos.
KernelEvolve aborda este desafío fundamental de la computación distribuida al transformar la optimización de kernels en un problema de búsqueda automatizado. Al tratar la generación de código como una exploración iterativa en un espacio de soluciones, en lugar de una tarea de generación única, el sistema puede adaptarse continuamente a nuevas plataformas y modelos, reduciendo drásticamente el tiempo de desarrollo y mejorando el rendimiento de inferencia y entrenamiento en entornos de producción masivos. Este enfoque es una evolución de los compiladores JIT y los auto-tuners, llevando la meta-optimización a un nivel agentic.
Arquitectura del Sistema
KernelEvolve se estructura como un sistema de IA agentic que formaliza la optimización de kernels como un problema de búsqueda. Sus componentes clave incluyen un LLM Synthesizer, un Tree Search Engine, una Retrieval-Augmented Knowledge Base y un Automated Evaluation Framework, todo ello orquestado por un 'job harness' de larga duración.
El LLM Synthesizer genera kernels candidatos en múltiples lenguajes (Triton, TLX, CuTe DSL, FlyDSL, CUDA, HIP, MTIA C++), utilizando prompts dinámicos y contextuales enriquecidos con diagnósticos de tiempo de ejecución y restricciones de hardware. El Tree Search Engine explora el espacio de optimización mediante algoritmos de búsqueda basados en grafos, como Monte Carlo tree search o estrategias evolutivas. Cada candidato es un nodo en el árbol de búsqueda, y el motor selecciona, transforma y evalúa, equilibrando la explotación de estrategias conocidas con la exploración de enfoques novedosos. Un mecanismo de memoria configurable permite a los nodos heredar contexto de sus padres o hermanos para refinar trayectorias de optimización o inyectar diversidad.
La Retrieval-Augmented Knowledge Base es jerárquica y contiene restricciones de corrección, guías de optimización agnósticas a la plataforma y documentación específica de hardware (manuales de arquitectura, conjuntos de instrucciones). Esta base de conocimiento es dinámica y se auto-evoluciona, destilando estrategias exitosas en 'skills' reutilizables que se retroalimentan, actuando como una forma de aprendizaje por refuerzo en contexto. Finalmente, el Automated Evaluation Framework valida la corrección y el rendimiento de cada kernel generado, utilizando un stack de herramientas de profiling (TritonBench, PyTorch Profiler, NCU, Proton, MTIA Insight) que proporcionan métricas a nivel de hardware. Este framework unifica el análisis a través de una abstracción centrada en el compilador, insertando instrumentación a nivel de MLIR y sintetizando trazas estructuradas para retroalimentar al LLM Synthesizer con diagnósticos detallados (ej. si el cuello de botella es memory-bound o compute-bound).
Flujo de Optimización de Kernel con KernelEvolve
- 1 Ingeniero Especifica operador, hardware y objetivos de rendimiento
- 2 Knowledge Base Recupera documentación y patrones de optimización
- 3 LLM Synthesizer Genera candidatos de kernel con prompts contextuales
- 4 Evaluation Framework Valida corrección y mide rendimiento (profiling detallado)
- 5 Tree Search Engine Analiza resultados, selecciona candidatos, aplica transformaciones
- 6 Loop Itera hasta alcanzar objetivo o agotar presupuesto
- 7 KernelEvolve Entrega kernel optimizado y validado para producción
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | NVIDIA GPUs, AMD GPUs, Meta MTIA, CPUs | Hardware de destino para la ejecución de kernels optimizados. |
| data-processing | LLM Synthesizer | Generación de código de kernel en múltiples DSLs y lenguajes de bajo nivel (Triton, CUDA, HIP, MTIA C++). Prompts dinámicos y contextuales, enriquecidos con diagnósticos. |
| orchestration | Tree Search Engine | Exploración del espacio de optimización de kernels utilizando algoritmos como Monte Carlo tree search y estrategias evolutivas. Mecanismo de memoria configurable para nodos del árbol de búsqueda. |
| storage | Retrieval-Augmented Knowledge Base | Almacenamiento y recuperación de documentación de hardware, guías de optimización y 'skills' auto-evolutivas. Jerárquica, con categorías para corrección, optimización agnóstica y específica de hardware. |
| observability | Automated Evaluation Framework (TritonBench, PyTorch Profiler, NCU, Proton, MTIA Insight) | Validación de corrección y profiling detallado del rendimiento de los kernels generados. Unificación de herramientas de profiling a través de una abstracción centrada en el compilador (MLIR). |
Trade-offs
Ganancias
- ▲▲ Development Velocity
- ▲ Inference Throughput
- ▲ Training Throughput
- ▲ Hardware Coverage
Costes
Fundamentos Teóricos
El problema de la optimización de código para arquitecturas de hardware específicas es un tema recurrente en la informática, con raíces en la teoría de compiladores y la optimización de programas. La idea de tratar la optimización como un problema de búsqueda se remonta a trabajos tempranos en compiladores auto-optimizantes y sistemas de tuning automático, como el trabajo de ATLAS (Automatically Tuned Linear Algebra Software) de Whaley et al. (1998) para BLAS, que utilizaba búsqueda empírica para encontrar las mejores implementaciones de kernels de álgebra lineal.
La aproximación de KernelEvolve, que combina LLMs con búsqueda en árboles y aprendizaje por refuerzo, puede verse como una evolución de estos principios, aplicando técnicas modernas de IA a un problema clásico. La integración de una base de conocimiento aumentada por recuperación (RAG) para inyectar información específica de hardware no vista durante el entrenamiento del LLM, se alinea con los avances en el procesamiento del lenguaje natural para mejorar la factualidad y la capacidad de razonamiento de los modelos. El concepto de 'self-evolving skill library' y el uso de trayectorias agentic para post-entrenar modelos especializados resuenan con los principios del aprendizaje por refuerzo y el meta-aprendizaje, donde un sistema mejora su capacidad de aprender y optimizar con la experiencia, similar a los trabajos de Schmidhuber sobre meta-aprendizaje y redes neuronales recurrentes.