Resumen Técnico: AutoMegaKernel (AMK)

Declaración del Problema

La decodificación autorregresiva de flujo único (single-stream) con un tamaño de lote (batch size) de 1 está fundamentalmente limitada por el ancho de banda. El piso de latencia está determinado por la relación entre los bytes totales de peso y el ancho de banda de la HBM ($t_{min} = \text{bytes\_de\_peso} / \text{ancho\_de\_banda\_HBM}$). Los marcos de ejecución convencionales (por ejemplo, PyTorch, cuBLAS) no logran acercarse a este piso porque lanzan un kernel separado para cada operador, incurriendo en una latencia de lanzamiento de CPU significativa y forzando a las activaciones a realizar viajes de ida y vuelta a la HBM entre operaciones. Aunque los CUDA Graphs amortizan la sobrecarga de lanzamiento, mantienen los límites inter-op y los viajes de ida y vuelta a la HBM.

Los enfoques de "megakernel" existentes (por ejemplo, Mirage Persistent Kernel) o los kernels optimados a mano abordan el piso de latencia fusionando todo el paso hacia adelante (forward pass) en un único kernel persistente cooperativo. Sin embargo, estos sistemas carecen de confianza y portabilidad: a menudo omiten puertas de seguridad estáticas (riesgo de bloqueos o condiciones de carrera), se dirigen a modelos/arquitecturas específicos de forma manual y no proporcionan una interfaz estructurada para que agentes autónomos optimicen el sistema de forma segura.

Metodología

AutoMegaKernel (AMK) es un sistema autónomo impulsado por agentes que compila modelos de la familia Llama de HuggingFace en un único kernel persistente cooperativo sin necesidad de escribir CUDA a mano para cada modelo. El sistema se estructura en cuatro capas y dos bucles de investigación automática (autoresearch loops):

1. Capa 0 (La VM): Un megakernel persistente escrito a mano, congelado y por arquitectura (vm/scheduler.cu). Lanza un threadblock cooperativo por cada Streaming Multiprocessor (SM). Cada bloque posee una cola local y recorre un orden topológico global de tareas, ejecutando instrucciones wait → dispatch → signal. Se basa en contadores monotónicos para la sincronización, evitando bloqueos (locks).
2. Capa 1 (Instrucciones): Micro-kernels que cumplen con el ABI (GEMV, atención, RMSNorm, etc.) y que son verificados en cuanto a corrección de forma aislada.
3. Capa 2 (Programador/Scheduler): La superficie de edición para el agente. Importa un modelo, lo reduce a un IR de Programación tipado (un DAG de tareas sincronizado por contadores monotónicos) y propone configuraciones estructurales (tiling, fusión, asignación de SM, profundidad de pipelining) mediante un objeto ScheduleConfig.
4. Capa 3 (Roadmap): Espacio reservado para la futura dinamismo (por ejemplo, continuous batching).

La Puerta de Seguridad (The Safety Gate):
La innovación central es un validador verificado estáticamente (validate() en schedule/ir.py). Antes de cualquier lanzamiento en la GPU, el validador realiza un conjunto de comprobaciones de grafo estáticas sobre el IR de la Programación para certificar:

• Ausencia de Deadlock: Asegurar que el DAG de productor-consumidor sea acíclico, que los umbrales de espera sean satisfactorios ($1 \le t \le \#productores$) y que los órdenes de las colas por SM sean válidos.
• Ausencia de Race (Condiciones de Carrera): Asegurar que los contadores compartidos se unan correctamente (todos los productores deben terminar antes de que un consumidor espere en un contador de múltiples productores) y que los accesos a memoria respeten las relaciones de "happens-before".
• Corrección: Verificar la alcanzabilidad de la salida y la integridad referencial.

Si un programa falla estas comprobaciones, es rechazado inmediatamente. Esto evita que programas propuestos por agentes de forma insegura causen bloqueos en la GPU.

Bucle de Investigación Automática (Autoresearch Loop):
AMK incluye un bucle no supervisado y controlable por agentes que propone iterativamente cambios en el ScheduleConfig, valida dichos cambios, ejecuta puertas de corrección contra una referencia en CPU, mide la latencia, y mantiene o revierte los cambios. Este bucle mejora autónomamente el megakernel sobre su propia línea base.

Contribuciones Clave

1. IR de Programación Verificado Estáticamente: Un validador congelado que certifica la ausencia de deadlocks y condiciones de carrera mediante comprobaciones de grafos estáticos (no pruebas mecanizadas) sobre una base de confianza. Logró cero falsos positivos en 7,160 programas adversarios construidos (6,091 inseguros) mientras aceptaba los 360 reducciones de modelos reales.
2. Cobertura de Generación Automática: AMK genera automáticamente megakernels correctos para 10 de los 10 modelos soportados (incluyendo checkpoints reales como SmolLM2-135M y TinyLlama-1.1B) sin escribir CUDA a mano. Rechaza con éxito variantes no soportadas (por ejemplo, proyecciones sesgadas) o documenta sus puntos ciegos con honestidad.
3. Retargeting Autónomo: Un único código fuente genera megakernels correctos para sm_80 (A100), sm_90 (H100) y sm_120 (RTX 5090), igualando la ejecución eager bit a bit o dentro de tolerancias estrictas de fp32 ($\le 3.8 \times 10^{-5}$).
4. Megakernels Cuantizados: El sistema genera automáticamente kernels de pesos cuantizados int8 (pérdida de precisión mínima/greedy-lossless, 1.12× más rápido) e int4 (con pérdida/lossy, 2.24× reducción de tráfico de pesos) desde la misma ruta de IR.
5. Compilador Controlable por Agentes: El sistema trata la corrección como un invariante estructural, permitiendo a un agente buscar en el espacio de programación de forma segura sin riesgo de colgar la GPU.
6. Estudio de Rendimiento Honesto: El artículo reporta resultados en GPUs de clase inferencia y de clase entrenamiento, revelando explícitamente dónde el sistema gana y dónde queda rezagado frente a las líneas base establecidas.

Resultados Experimentales

Corrección y Portabilidad:

• En la RTX 5090, A100 y H100, AMK reproduce la decodificación greedy de HuggingFace token por token y coincide con la perplejidad de teacher-forcing hasta $2.5 \times 10^{-7}$ en el checkpoint de SmolLM2-135M.
• El validador rechazó todos los 6,091 programas inseguros en la prueba de estrés, demostrando una tasa de falsos positivos del 0%.

Rendimiento en GPUs de Clase Inferencia (La Victoria):

• RTX 5090 (Consumo): Un megakernel auto-optimizado de solo peso int8 (W8A16) supera a cuBLAS con CUDA-graph (bf16) por 1.19–1.23× en decodificación de batch-1. Esta es una comparación de asimetría de precisión (int8 vs. bf16) donde la victoria proviene de transmitir menos bytes de peso (0.61× los bytes), no de un kernel más rápido por byte.
• Flota de Inferencia de Datacenter: El megakernel int8 supera a cuBLAS con CUDA-graph bf16 en toda la flota de clase inferencia:
  • NVIDIA L4: Hasta 1.33× más rápido (crece con el tamaño del modelo).
  • L40S: 1.25–1.27× más rápido.
  • A10G: Hasta 1.08× más rápido a escala.
• Auto-mejora: El bucle de investigación automática mejoró autónomamente el megakernel sobre su valor por defecto en un 1.25–1.72× en la RTX 5090.

Rendimiento en GPUs de Clase Entrenamiento (La Derrota):

• En GPUs de alto ancho de banda de clase entrenamiento (A100, H100), el megakernel generado automáticamente se queda atrás de cuBLAS con CUDA-graph y vLLM.
  • En H100, AMK fue 1.65× más lento que la ruta de CUDA-graph por defecto de vLLM.
  • En A100, el GEMV optimizado en bf16 alcanzó solo aproximadamente el 63% del pico de HBM medido (frente al ~90% de cuBLAS).
• Diagnóstico: El cuello de botella es el costo de sincronización entre SM por tile (un costo fijo que el megakernel paga y cuBLAS no). En silicio de alto ancho de banda, este costo de sincronización domina, mientras que en GPUs de clase inferencia, los ahorros de ancho de banda mediante la fusión y la cuantización amortizan este costo.

Significancia y Reivindicaciones

El artículo afirma que AutoMegaKernel representa un cambio de la "velocidad bruta" hacia la seguridad sistemática y la portabilidad en la síntesis de megakernels. Su principal contribución es una arquitectura de compilador verificada estáticamente que permite a los agentes autónomos explorar de forma segura el espacio de los programas de modelos completos.

• Confianza: Reemplaza el riesgo de colgar las GPUs con una puerta de seguridad estática que rechaza programas inseguros antes del lanzamiento.
• Generalidad: Logra un retargeting autónomo a través de arquitecturas (sm_80/90/120) y familias de modelos sin necesidad de ajuste manual de CUDA.
• Honestidad: Los autores enmarcan explícitamente los resultados de rendimiento según el régimen de uso. El sistema gana en GPUs de clase inferencia al fusionar operaciones y aprovechar la cuantificación para reducir el tráfico de memoria, pero actualmente pierde en GPUs de clase entrenamiento debido a los costos de sincronización que requieren una programación de grano más grueso (trabajo futuro) en lugar de simplemente mejores micro-kernels.

El artículo concluye que, si bien la calidad actual de los kernels aún no iguala a las líneas base optimadas a mano en hardware de alto rendimiento para entrenamiento, la estructura de soporte verificada estáticamente proporciona la base necesaria para automatizar de forma segura las mejoras futuras (por ejemplo, sincronización de grano más grueso, pipelining más profundo) necesarias para cerrar esa brecha.