Resumen Técnico: GPU Forecasters: Modelos de Lenguaje como Sustitutos Selectivos para la Optimización del Tiempo de Ejecución de Kernels

Declaración del Problema

Los kernels de GPU son las herramientas fundamentales de la computación profunda moderna, donde las mejoras menores en la eficiencia se traducan en reducciones significativas en los costos de entrenamiento e inferencia. La optimización de estos kernels implica típicamente un ciclo de generación-evaluación en el que los Modelos de Lenguaje de Gran Escala (LLMs) proponen kernels candidatos, que luego son compilados y ejecutados en hardware físico para medir su rendimiento. Sin embargo, este proceso de evaluación es un cuello de botella severo. Cada evaluación requiere compilación (que toma entre 60 y 120 segundos para CUDA) y ejecución en el dispositivo, lo que hace que el rendimiento de la evaluación sea órdenes de magnitud más lento que la generación. A medida que los presupuestos de búsqueda escalan, el costo de las mediciones físicas repetidas limita la capacidad de explorar el vasto espacio de implementaciones potenciales de kernels.

Mientras que trabajos previos se han centrado en usar LLMs para generar mejores kernels, este artículo aborda el problema complementario de reducir el costo de evaluación. Los autores proponen utilizar los LLMs no solo como generadores, sino como sustitutos selectivos (selective surrogates) de la GPU física. El objetivo es pronosticar el rendimiento relativo de un kernel candidato frente a una implementación de referencia e identificar cuándo el modelo tiene incertidumbre, de modo que solo los casos inciertos se difieran hacia la costosa GPU física.

Metodología

Formulación de la Tarea: Pronóstico Selectivo

Los autores formulan el problema como el pronóstico de la aceleración relativa ($S_i = T_{ref} / T_i$) de un kernel candidato comparado con un kernel de referencia, dada la especificación del hardware. Predecir el tiempo de ejecución absoluto en segundos se considera intratable debido a la complejidad de modelar el comportamiento del compilador y los detalles microarquitectónicos. En su lugar, el objetivo se discretiza en ocho contenedores (bins) ordinales de aceleración (que van desde "Ralentización Severa" hasta "Aceleración Extrema").

El sustituto (un LLM) toma el kernel de referencia, el kernel candidato y las especificaciones del hardware como entrada. Genera una cadena de razonamiento y emite una distribución de probabilidad sobre los ocho contenedores. La confianza declarada del modelo se define como la probabilidad máxima asignada a cualquier contenedor individual.

Estrategia de Predicción Selectiva

El sistema opera bajo un marco de predicción selectiva. El sustituto se utiliza para clasificar candidatos. Si la confianza del sustituto supera cierto umbral, el pronóstico se acepta y el candidato no se mide en la GPU. Si la confianza es baja (indicando incertidumbre), el candidato se difiere a la GPU física para la medición de la verdad de campo (ground-truth). Esto permite que el proceso de búsqueda considere muchos más candidatos dentro de un presupuesto fijo de medición de GPU.

Aprendizaje por Refuerzo (RL) para la Calibración

Para mejorar la precisión del sustituto y, crucialmente, su calibración (la alineación entre la confianza declarada y la precisión real), los autores aplican Aprendizaje por Refuerzo a un modelo de pesos abiertos (GPT-OSS-20B). Definen recompensas basadas en:

1. Corrección: Si el contenedor predicho coincide con el contenedor medido.
2. Calibración Categórica (Brier Score): Penalizar la distribución basándose en la distancia respecto al objetivo real de un solo bit (one-hot).
3. Calibración Sensible a la Distancia (CRPS): Penalizar la distribución acumulada, teniendo en cuenta la naturaleza ordinal de los contenedores de aceleración (por ejemplo, predecir una aceleración de 2× cuando la realidad es de 1.5× es menos grave que predecir 4×).

Los datos de entrenamiento se construyen como un subproducto de la propia búsqueda de kernels, utilizando candidatos exitosos de ejecuciones de búsqueda previas como ejemplos de entrenamiento.

Contribuciones Clave

1. LLMs como Sustitutos de GPU: El artículo demuestra que los LLMs pueden lograr una precisión no trivial al pronosticar cómo responde una GPU física a los kernels candidatos, recuperando una fracción sustancial de la aceleración disponible al clasificar candidatos basándose en pronósticos en lugar de mediciones.
2. Marco de Predicción Selectiva: Los autores presentan la predicción de rendimiento de GPU como un problema de predicción selectiva, demostiendo que los sustitutos pueden reportar incertidumbre calibrada. Esto permite una estrategia de búsqueda que confía en los pronósticos de alta confianza y reserva las mediciones físicas para los casos inciertos.
3. Impacto del RL en el Comportamiento del Sustituto: A través de experimentos, el artículo caracteriza cómo el entrenamiento por RL desplaza las propiedades del sustituto. Específicamente, añadir recompensas de calibración (Brier o CRPS) mejora la fiabilidad de las estimaciones de confianza del modelo, mientras que entrenar solo en corrección puede degradar la relación entre confianza y precisión.
4. Mejora de la Búsqueda de Extremo a Extremo: El estudio muestra que integrar el sustituto en un bucle de búsqueda de kernels (específicamente una búsqueda basada en PUCT) permite que la búsqueda explore significativamente más candidatos por cada medición de GPU. En la mayoría de las tareas, esto conduce al hallazgo de kernels más rápidos que un baseline que mide cada propuesta bajo el mismo presupuesto de GPU.
5. Lanzamiento de Dataset: Los autores lanzan una colección reutilizable de 12,388 kernels de GPU generados por LLM con tiempos de ejecución medidos, que abarcan 118 problemas, múltiples tipos de GPU (A100, L40S, H100) y diversos métodos de generación. Este dataset soporta una investigación más amplia en predicción de rendimiento de GPU y búsqueda de kernels.

Resultados

• Precisión del Pronóstico: Los LLMs estándar (ej. GPT-OSS-20B) recuperan aproximadamente el 83% de la mejor aceleración disponible cuando solo se mide el 1% superior de los candidatos en la GPU. Este rendimiento mejora al 94% cuando se mide el 50% superior.
• Calibración y RL: El entrenamiento por RL altera significativamente el comportamiento del sustituto.
  • El entrenamiento basado solo en corrección mejora la precisión del contenedor exacto pero degrada la correlación entre confianza y error (el modelo se vuelve excesivamente confiado en predicciones erróneas).
  • El entrenamiento de calibración (Brier/CRPS) mejora la monotonicidad del error de pronóstico frente a la confianza. Notablemente, el entrenamiento calibrado con Brier mejora la métrica de "aceleración recuperada" en aproximadamente 2.1 puntos porcentuales respecto al modelo base, lo que sugiere que las distribuciones mejor calibradas conducen a mejores señales de clasificación para el algoritmo de búsqueda.
• Eficiencia de Búsqueda: En la búsqueda de kernels de extremo a extremo, la búsqueda habilitada por el sustituto encuentra kernels más rápidos que el baseline en cuatro de seis tareas, alcanzando a menudo la mejor aceleración del baseline con sustancialmente menos mediciones de GPU. En las tareas donde el baseline gana, la búsqueda parece saturarse temprano, indicando que el sustituto es más beneficioso cuando el espacio de búsqueda es grande y no explorado.
• Momentos de Descubrimiento: El sustituto puede identificar "momentos de descubrimiento" (mutaciones donde un kernel hijo es significativamente más rápido que su padre) mejor que el azar, aunque la precisión sigue siendo un desafío, lo que requiere confirmación por GPU para decisiones de alto riesgo.

Significado y Reivindicaciones

El artículo sostiene que los LLMs pueden desempeñar un papel más amplio en la optimización de kernels más allá de la mera generación de código. Al actuar como modelos virtuales de una GPU, pueden servir como filtros de alto rendimiento y bajo costo que gestionan la escasez de recursos de hardware físico.

Los autores enfatizan que el valor principal de este enfoque no es solo la velocidad, sino la capacidad de realizar una optimización basada en modelos sin consultar el sistema original en cada paso. El sustituto ofrece un mayor rendimiento y un menor costo por consulta que una GPU, intercambiando el ruido del pronóstico que puede gestionarse mediante estimaciones de incertidumbre calibradas. El trabajo conecta la optimización de kernels con el campo más amplio de la construcción de simuladores neuronales para sistemas físicos complejos, sugiriendo un camino hacia optimizadores que traten la medición de hardware como un recurso escaso para ser utilizado selectivamente.