Escribir código GPU rápido es una de las especializaciones más agotadoras en ingeniería de aprendizaje automático. Los investigadores de RightNow AI quieren automatizarlo por completo.
El equipo de investigación de RightNow AI ha lanzado AutoKernel, un marco de código abierto que aplica un bucle de agente LLM autónomo a la optimización del kernel de GPU para modelos PyTorch arbitrarios. El enfoque es sencillo: dale cualquier modelo antes de irte a dormir y despierta con núcleos Triton más rápidos, sin necesidad de tener experiencia en GPU.
Por qué los kernels de GPU son tan difíciles de optimizar
Un kernel de GPU es una función que se ejecuta en paralelo en miles de núcleos de GPU. Cuando ejecuta un modelo de transformador como LLaMA o GPT-2, la mayor parte del tiempo de cómputo se gasta dentro de los núcleos para operaciones como multiplicación de matrices (matmul), softmax, normalización de capas y atención. Estos núcleos viven en bibliotecas como cuBLAS y cuDNN, o se generan automáticamente mediante el proceso de compilación de PyTorch.
El problema es que para obtener el máximo rendimiento de estos núcleos se requiere razonar simultáneamente sobre la intensidad aritmética, la fusión de la memoria, la presión de registro, los tamaños de los mosaicos, la sincronización del nivel de deformación y la selección de instrucciones del núcleo del tensor: una combinación de habilidades que lleva años desarrollar. Un único núcleo matmul de alto rendimiento puede implicar más de 200 líneas de código CUDA o Triton con docenas de parámetros interdependientes. Esta experiencia es escasa y el proceso de ajuste manual escala mal a medida que evolucionan las arquitecturas de los modelos.
El conjunto de pruebas KernelBench, que evalúa los LLM de vanguardia en 250 problemas del kernel de GPU, descubrió que incluso los mejores modelos igualaban el rendimiento básico de PyTorch en menos del 20% de los casos utilizando la generación de una sola vez. AutoKernel se creó directamente en respuesta a esa brecha.
El bucle: editar, comparar, mantener o revertir
La idea central de AutoKernel es que el flujo de trabajo de un ingeniero de kernel experto es en sí mismo un bucle simple: escribir un candidato, compararlo, mantener las mejoras, descartar las regresiones, repetir. El marco mecaniza este bucle. Un agente LLM modifica un solo archivo (kernel.py); un sistema de referencia fijo verifica la corrección y mide el rendimiento, y el resultado determina si el cambio persiste. Fundamentalmente, cada experimento se asigna a un compromiso de git. Los experimentos mantenidos hacen avanzar la rama; Los experimentos revertidos se borran limpiamente con git reset. Se puede navegar por todo el historial con herramientas estándar de Git, y los resultados del experimento se registran en un archivo results.tsv separado por tabulaciones: libre de dependencias, legible por humanos y fácilmente analizable por el agente.
Cada iteración toma aproximadamente 90 segundos: 30 segundos para verificar la corrección, 30 segundos para evaluar el rendimiento a través de do_bench de Triton y 30 segundos para el razonamiento del agente y la modificación del código. A aproximadamente 40 experimentos por hora, una ejecución de 10 horas durante la noche produce entre 300 y 400 experimentos en varios núcleos.
Este diseño se basa directamente en el proyecto de investigación automática de Andrej Karpathy, que demostró que un agente de IA que ejecutara un ciclo de mantenimiento/reversión en el código de entrenamiento LLM podría descubrir 20 optimizaciones en 700 experimentos en dos días en una sola GPU. AutoKernel trasplanta este bucle al código del kernel, con un espacio de búsqueda diferente y un punto de referencia controlado por la corrección como función de evaluación en lugar de pérdida de validación.
El agente lee un documento de instrucciones de 909 líneas llamado program.md, que codifica el conocimiento experto en un manual de optimización de seis niveles. Los niveles progresan desde el ajuste del tamaño del bloque (barrer las dimensiones de los mosaicos a través de potencias de 2, ajustar num_warps y num_stages) a través de patrones de acceso a la memoria (cargas fusionadas, captación previa de software, swizzling L2), optimizaciones informáticas (acumulación de TF32, fusión de epílogo), técnicas avanzadas (split-K, núcleos persistentes, autoajuste de Triton, especialización de deformación), estrategias específicas de la arquitectura (TMA en Hopper, cp.async en Ampere, tamaños ajustados para L4/RTX) y, finalmente, algoritmos específicos del kernel como softmax en línea para atención y el algoritmo de Welford para normalización. El documento de instrucciones es intencionalmente completo para que el agente pueda funcionar más de 10 horas sin quedarse atascado.
Primero, crear perfiles y optimizar lo que importa
A diferencia de trabajos anteriores que tratan los problemas del kernel de forma aislada, AutoKernel parte de un modelo PyTorch completo. Utiliza torch.profiler con registro de forma para capturar el tiempo de GPU por kernel, luego clasifica los objetivos de optimización usando la ley de Amdahl: el principio matemático de que la aceleración general que puedes lograr está limitada por la cantidad de tiempo de ejecución total que representa ese componente. Una aceleración de 1,5 veces en un kernel que consume el 60% del tiempo de ejecución total produce una ganancia de extremo a extremo de 1,25 veces. La misma aceleración en un kernel que consume el 5% del tiempo de ejecución produce sólo 1,03×.
El generador de perfiles detecta el hardware de GPU a partir de una base de datos de especificaciones conocidas que cubren los aceleradores NVIDIA (H100, A100, L40S, L4, A10, RTX 4090/4080/3090/3080) y AMD (MI300X, MI325X, MI350X, MI355X). Para GPU desconocidas, estima el rendimiento máximo de FP16 a partir del recuento de SM, la velocidad de reloj y la capacidad de cómputo, lo que hace que el sistema se pueda utilizar en una gama más amplia de hardware que solo las últimas ofertas de NVIDIA.
El orquestador (orchestrate.py) pasa de un kernel al siguiente cuando se cumple cualquiera de las cuatro condiciones: cinco reversiones consecutivas, se alcanzó el 90 % de la utilización máxima de la GPU, un presupuesto de tiempo transcurrido de dos horas o una aceleración de 2 veces ya lograda en ese kernel. Esto evita que el agente dedique demasiado tiempo a núcleos con rendimientos decrecientes mientras esperan los objetivos de mayor impacto.
Arnés de corrección de cinco etapas
El rendimiento sin corrección es inútil y AutoKernel es particularmente minucioso en este frente. Cada núcleo candidato pasa por cinco etapas de validación antes de que se registre cualquier aceleración. La etapa 1 ejecuta una prueba de humo en una pequeña entrada para detectar errores de compilación y discrepancias de forma en menos de un segundo. La etapa 2 abarca de 8 a 10 configuraciones de entrada y tres tipos de datos (FP16, BF16 y FP32) para detectar errores que dependen del tamaño, como el manejo de límites y la lógica del resto de mosaicos. La etapa 3 prueba la estabilidad numérica bajo entradas adversas: para softmax, filas de grandes valores idénticos; para matmul, rango dinámico extremo; para normalización, varianza cercana a cero. La etapa 4 verifica el determinismo ejecutando la misma entrada tres veces y requiriendo salidas idénticas bit a bit, lo que detecta las condiciones de carrera en reducciones paralelas y atómicas no deterministas. La etapa 5 prueba dimensiones sin potencia de dos como 1023, 4097 y 1537 para exponer errores de enmascaramiento y errores de resto de mosaicos.
Las tolerancias son específicas del tipo: FP16 usa atol = 10⁻², BF16 usa 2 × 10⁻² y FP32 usa 10⁻⁴. En la evaluación completa del artículo en 34 configuraciones en una NVIDIA H100, las 34 aprobaron la corrección sin fallas en las salidas del kernel entusiastas, compiladas y personalizadas.
Backend dual: Triton y CUDA C++
AutoKernel admite backends Triton y CUDA C++ dentro del mismo marco. Triton es un lenguaje de dominio específico similar a Python que compila JIT en 1 a 5 segundos, lo que lo hace ideal para una iteración rápida: el agente puede modificar el tamaño de los bloques, los conteos de deformación, las etapas de la canalización, la precisión del acumulador y la estructura del bucle. Triton alcanza habitualmente entre el 80 y el 95 % del rendimiento de cuBLAS para matmul. CUDA C++ se incluye para casos que requieren acceso directo a primitivas de nivel warp, instrucciones de núcleo tensor WMMA (usando fragmentos de 16×16×16), cargas vectorizadas a través de float4 y half2, diseños de memoria compartida sin conflictos de bancos y doble almacenamiento en búfer. Ambos backends exponen la misma interfaz kernel_fn(), por lo que la infraestructura de referencia se ejecuta de manera idéntica independientemente del backend.
El sistema cubre nueve tipos de kernel que abarcan las operaciones dominantes en las arquitecturas de transformadores modernas: matmul, flash_attention, fusionado_mlp, softmax, Layernorm, rmsnorm, cross_entropy, Rotary_embedding y reduce. Cada uno tiene una implementación de referencia de PyTorch en reference.py que sirve como oráculo de corrección, y el punto de referencia calcula el rendimiento en TFLOPS o GB/s junto con la utilización de la línea del techo frente al pico de GPU detectado.
Resultados de referencia en H100
Medidos en una GPU NVIDIA H100 HBM3 de 80 GB (132 SM, capacidad de procesamiento 9.0, CUDA 12.8) frente a PyTorch ansioso y torch.compile con max-autotune, los resultados para los núcleos vinculados a la memoria son significativos. RMSNorm logra 5,29 veces más de ansioso y 2,83 veces más de torch.compile en el tamaño más grande probado, alcanzando 2788 GB/s, el 83 % del ancho de banda máximo de 3352 GB/s del H100. Softmax alcanza los 2.800 GB/s con una aceleración de 2,82 veces superior a la de ansioso y de 3,44 veces superior a la de torch.compile. La entropía cruzada alcanza 2,21 veces más que la velocidad y 2,94 veces más que la antorcha.compilación, alcanzando 2070 GB/s. Las ganancias de estos núcleos provienen de la fusión de descomposiciones ATen de múltiples operaciones en núcleos Triton de un solo paso que minimizan el tráfico HBM (memoria de alto ancho de banda).
AutoKernel supera a torch.compile en 12 de las 16 configuraciones representativas evaluadas en el documento, a pesar de que torch.compile con max-autotune ejecuta su propio autoajuste Triton. La fusión y el autoajuste genéricos de TorchInductor no siempre encuentran las estrategias especializadas de mosaico y reducción que explotan las implementaciones específicas del kernel.
Matmul es notablemente más difícil: el backend cuBLAS de PyTorch está ampliamente ajustado según la arquitectura de GPU. El iniciador Triton alcanza los 278 TFLOPS, muy por debajo de cuBLAS. Sin embargo, en el tamaño de 2048³, AutoKernel supera a torch.compile por 1,55×, lo que demuestra que el autoajuste matmul de TorchInductor tampoco siempre es óptimo. Cerrar la brecha de cuBLAS sigue siendo el objetivo principal para la iteración continua del agente.
En la implementación comunitaria, un kernel optimizado para AutoKernel ocupó el primer lugar en la clasificación de vectoresum_v2 B200 con una latencia de 44.086 µs, superando al segundo lugar con 44.249 µs y al tercer lugar con 46.553 µs. Un usuario de la comunidad también informó que un solo mensaje de AutoKernel, que requirió aproximadamente tres minutos de interacción del agente, produjo un núcleo de multiplicación de matrices Triton FP4 que supera a CUTLASS entre 1,63 y 2,15 veces en múltiples formas en H100. CUTLASS representa un código de plantilla C++ optimizado manualmente diseñado específicamente para núcleos tensoriales NVIDIA, lo que hace que este resultado sea particularmente notable.
Conclusiones clave
AutoKernel convierte semanas de ajuste experto de GPU en un proceso autónomo de la noche a la mañana. Al mecanizar el ciclo de escritura, mantenimiento y reversión del punto de referencia que ya siguen los ingenieros expertos en kernel, el sistema ejecuta de 300 a 400 experimentos por sesión nocturna en una sola GPU sin ninguna intervención humana. La corrección no es negociable antes de que se registre cualquier aceleración. Cada núcleo candidato debe pasar un arnés de cinco etapas que cubre pruebas de humo, barridos de forma en más de 10 configuraciones, estabilidad numérica bajo entradas adversas, verificación de determinismo y casos extremos sin potencia de dos, eliminando el riesgo de que el agente “optimice” su camino hacia resultados incorrectos. Los núcleos vinculados a la memoria obtienen las mayores ganancias sobre PyTorch ansioso y torch.compile. En una NVIDIA H100, los núcleos Triton de AutoKernel logran 5,29 veces más de entusiasmo en RMSNorm, 2,82 veces en softmax y 2,21 veces en entropía cruzada, y las ganancias provienen de la fusión de descomposiciones ATen de operaciones múltiples en núcleos de un solo paso que minimizan el tráfico de HBM. La ley de Amdahl determina dónde pasa su tiempo el agente. En lugar de optimizar los kernels de forma aislada, AutoKernel perfila todo el modelo PyTorch y asigna el esfuerzo proporcionalmente a la participación de cada kernel en el tiempo de ejecución total de la GPU, lo que garantiza que las mejoras se acumulen a nivel de modelo, no solo a nivel de kernel.
Consulte el documento y el repositorio. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros