Compilar LLM en un MegaKernel para lograr inferencia de baja latencia
(zhihaojia.medium.com)- Se desarrolló un compilador que convierte automáticamente la inferencia de LLM en un único megakernel
- El enfoque MegaKernel (kernel persistente) integra por completo el cómputo y la comunicación de la inferencia de LLM en un solo kernel de GPU, lo que permite una latencia muy baja
- Existe el problema de que, debido a la estructura distribuida de los frameworks de ML y las bibliotecas de kernels existentes, la unificación de todo el pipeline en un solo kernel es muy difícil
- Mirage Persistent Kernel (MPK) convierte automáticamente la inferencia de LLM en múltiples GPU en un megakernel de alto rendimiento mediante un compilador y un sistema de runtime
- MPK transforma el grafo de operaciones en un grafo de tareas detallado, maximizando el software pipelining y la superposición entre cómputo y comunicación
- Al aplicar MPK, se reduce la latencia de generación de tokens frente a los sistemas existentes, y la mejora de rendimiento aumenta aún más conforme crece el número de GPU
Panorama general y ventajas del enfoque MegaKernel
- En la inferencia de modelos de lenguaje grandes (LLM), una de las formas más efectivas de reducir la latencia es fusionar todo el proceso de cómputo y comunicación en un único megakernel (kernel unificado)
- Este enfoque hace que un solo kernel de GPU ejecute sin interrupciones todo el procesamiento, incluidas las operaciones por capa de todo el modelo y la comunicación entre GPU
- Sus principales ventajas son las siguientes
- Omitir lanzamientos repetitivos de kernels para eliminar el overhead de lanzamiento de kernels
- Hacer posible el software pipelining a lo largo de las capas
- Ejecutar cómputo y comunicación al mismo tiempo para ocultar la latencia
Limitaciones previas y aparición de MPK
- Los frameworks de ML existentes como PyTorch, Triton y TVM no ofrecen de forma nativa la generación automática de un megakernel end-to-end
- Los sistemas LLM reales están compuestos por combinaciones de distintas bibliotecas de kernels como NCCL/NVSHMEM (comunicación), FlashInfer/FlashAttention (attention), CUDA/Triton (operaciones personalizadas), por lo que es difícil integrarlos en un solo kernel
- En este contexto, investigadores de CMU, UW, Berkeley, NVIDIA y Tsinghua desarrollaron Mirage Persistent Kernel (MPK)
- MPK combina compilador y runtime para convertir automáticamente todo el pipeline de inferencia de LLM en un megakernel de alto rendimiento
Valor central de MPK
- MPK elimina por completo el overhead de lanzamiento de kernels y maximiza la superposición entre operaciones, carga de datos y comunicación entre capas para habilitar inferencia de LLM de latencia ultrabaja
- En pruebas reales (prompt de 39 tokens, generación de 512 tokens, sin speculative decoding),
- En un entorno de una sola GPU NVIDIA A100 40GB, frente a la latencia de decodificación por token (14.5 ms) de sistemas optimizados existentes como vLLM/SGLang, MPK la reduce hasta 12.5 ms
- Esta cifra se acerca al límite inferior teórico (10 ms), considerando 1.6 TB/s de ancho de banda de memoria y carga de pesos de 16 GB
- Al integrar completamente cómputo y comunicación en entornos multi-GPU, la ventaja de rendimiento de MPK se vuelve más evidente a medida que aumenta el número de GPU
Estructura detallada de funcionamiento de MPK
Parte 1. Compilador – conversión del grafo de operaciones de LLM a grafo de tareas
- En general, las operaciones de un LLM se representan como un grafo computacional en el que cada operación (por ejemplo, multiplicación de matrices, attention) o operación de comunicación (por ejemplo, all-reduce) es un nodo, y las dependencias de datos son las aristas
- En los diseños existentes, es común ejecutar un kernel por operador, pero esto solo refleja dependencias a nivel kernel y no a nivel de los datos realmente dependientes, por lo que las oportunidades de pipelining son limitadas
- Ejemplo: si hay un all-reduce después de una multiplicación de matrices, el all-reduce solo empieza cuando termina toda la multiplicación. En realidad, es posible dividir los datos y aprovechar ejecuciones parciales y relaciones de dependencia más finas
- El compilador de MPK refina el grafo de operaciones y lo transforma automáticamente en un fine-grained task graph adecuado para unidades reales de datos
- Cada tarea (rectángulo) es una unidad de cómputo o comunicación asignada a un SM individual de GPU
- Cada evento (círculo) es un punto de sincronización entre tareas
- Las aristas entre tareas y eventos expresan de forma eficiente las dependencias de datos y control
- Gracias a este grafo de tareas, MPK permite que el cómputo y la comunicación se superpongan más, ya sea de forma parcial o en paralelo
- Con el Mirage kernel superoptimizer también se generan automáticamente implementaciones CUDA de alto rendimiento adecuadas para cada tarea
Parte 2. Runtime – ejecución del grafo de tareas dentro del megakernel
- El runtime de MPK ejecuta completamente el grafo de tareas solo dentro de un kernel de GPU (megakernel)
- Todos los SM (Streaming Multiprocessors) de la GPU se dividen estáticamente en roles de workers y schedulers
Workers
- Cada worker funciona a nivel de SM y administra una cola de tareas dedicada
- En un esquema de bucle
- toma la siguiente tarea de la cola
- la ejecuta (por ejemplo,
matmul,attention, transferencia de datos) - al completarla, notifica al evento correspondiente
- repite el proceso
- Esto permite optimizar el uso de recursos de cada worker y habilitar operaciones de capas asíncronas
Schedulers
- Un scheduler distribuido opera como un único warp dentro de cada SM, y se pueden ejecutar hasta 4 schedulers simultáneamente
- Cada scheduler administra una cola de eventos activados y asigna a los workers las tareas cuyas condiciones ya se cumplieron
- Así se hace posible la distribución masiva de tareas sin overhead de sincronización centralizada
Método de ejecución basado en eventos
- Cuando una tarea termina, incrementa un contador de evento específico. Cuando el contador alcanza un umbral, el evento se activa y se inserta en la cola del scheduler
- El scheduler ejecuta las tareas posteriores que dependen de ese evento
- Gracias a esto, el software pipelining de grano fino y la superposición entre cómputo y comunicación ocurren de forma natural
- Ejemplo: el
matmulde una capa y elattentionde otra capa pueden ejecutarse al mismo tiempo - En cuanto aparecen resultados parcialmente completados de
matmul, puede comenzar la comunicaciónall-reduce
- Ejemplo: el
- Como toda la planificación y el cambio entre tareas ocurre dentro del contexto de un solo kernel, el overhead entre tareas es muy bajo, del orden de 1–2 microsegundos (μs)
Dirección futura
-
Objetivo de MPK: permitir que los desarrolladores puedan compilar fácilmente LLM en un megakernel y obtener el máximo rendimiento escribiendo solo una pequeña cantidad de código Python (del orden de unas decenas de líneas)
-
Principales líneas de avance
- Soporte para arquitecturas de GPU recientes: por ejemplo, NVIDIA Blackwell y enfoques especializados a nivel warp
- Manejo de workloads dinámicos: investigación de estrategias de compilación para modelos que requieren control de flujo dinámico, como mixture-of-experts (MoE)
- Planificación avanzada de tareas: investigación y posible aplicación de políticas modernas como prioridades y optimización de throughput
-
MPK plantea un punto de inflexión fundamental en la forma de compilar y ejecutar cargas de trabajo de inferencia de LLM sobre GPU, y busca ampliar la colaboración con la comunidad
Material adicional
- El código y la documentación de MPK (Mirage Persistent Kernel), así como los resultados más recientes de investigación, pueden consultarse en GitHub (https://github.com/mirage-project/mirage)
1 comentarios
Opiniones de Hacker News
Al autor: me parece interesante que el enfoque de intérprete on-GPU se vea como una dirección futura muy prometedora. Hay otra investigación con una aproximación casi idéntica, así que recomiendo ver esta publicación relacionada. El modelo fundamental de programación de CUDA (por ejemplo, el lanzamiento de kernels) se está rodeando para lograr paralelización basada en tareas finas, y he visto directamente que este enfoque mejora más el aprovechamiento del hardware. Me pregunto si CUDA no nos ha estado limitando de varias maneras. También me entusiasma la posibilidad de que la investigación del autor llegue a convertirse en un backend experimental de PyTorch. Y un detalle menor: los dos párrafos de la primera parte son casi idénticos, así que lo señalo como una pequeña errata.
He trabajado de cerca por un tiempo con vLLM y SGLang, y estoy convencido de que este proyecto es exactamente la forma ideal de un proyecto de seguimiento. Me impresionó el análisis del grafo de dependencias de operaciones y la manera de fusionar operaciones o programar tareas de forma más inteligente. Felicitaciones al equipo.
Leí por encima la publicación y el README de GitHub, y me parece un proyecto realmente genial. Me pregunto si este tipo de optimización podría aplicarse no solo a inferencia, sino también a la etapa de entrenamiento. En particular, entiendo que fusionar la operación backward y la comunicación de gradientes sería un reto. Tengo entendido que por ahora no soporta cargas dinámicas de trabajo (por ejemplo, MoE), pero menciono un artículo reciente sobre procesar MoE en un solo kernel: FlashDMoE: Fast Distributed MoE in a Single Kernel.
Gracias por leer la publicación y hasta el README. También es posible dar soporte a la etapa de entrenamiento, pero como los kernels de entrenamiento suelen ser más grandes, el overhead de lanzamiento de kernels normalmente no es un problema tan importante, así que el mayor beneficiado es la inferencia, especialmente la de baja latencia. También vimos con interés el paper de FlashDMoE que compartiste, y destacamos que dar soporte a modelos MoE es nuestro próximo objetivo.
Personalmente, soy algo escéptico sobre dedicar tiempo a optimizar el entrenamiento basado en gradientes. En la práctica, muchas tareas de entrenamiento tienen características de valores discretos, así que creo que el aprendizaje basado en gradientes no las maneja bien.
El siguiente paso es compilar directo a Verilog y comprar hardware para LLM en AliExpress.
Comparten un texto que presenta tecnologías de hardware como Chisel. Antes de la aparición de la IA y las GPU, esta idea de pasar directamente del software al hardware era un enfoque prometedor. El avance de las CPU está estancado, y el deseo de optimizar más la capa intermedia entre software y hardware se ha mantenido, pero es probable que la computación paralela estilo GPU siga siendo el método principal de aceleración. Las CPU de propósito general probablemente terminarán quedándose como un pequeño cerebro que gestiona a la GPU. Aun así, se espera que el enfoque de pasar directamente del software al hardware tenga dificultades para volverse dominante.
Se estima que, si la estructura de los LLM se estabiliza en 5 a 10 años, mapearlos directamente al hardware podría volverse algo práctico. Se menciona que, con la tecnología actual, incluso modelos de decenas de miles de millones de parámetros podrían caber en una sola oblea usando solo compuertas lógicas de ultra baja precisión, cerca de 1.5 bits. A medida que aumenta la precisión, la cantidad de compuertas crece exponencialmente, así que por ahora resulta más eficiente mantener la memoria de pesos y compartir unidades de cómputo. En el futuro, desarrollar LLM de ultra baja precisión será una tarea esencial.
Broma sobre que, si el costo de entrenamiento ya es alto, agregar también el costo de la máscara complica todavía más las cosas, junto con la evaluación más fría de que, en realidad, los startups de hardware para IA llevan mucho tiempo intentando este tipo de dirección.
Si realmente existiera una solución tipo LLM-in-a-box, sería bastante atractiva. Pronto tendré la oportunidad de trabajar en un entorno offline (air-gap), y creo que una solución así sería muy útil.
Probé directamente el código en un entorno GPU de Modal, y las cifras de mejora de rendimiento que afirma la investigación sí se reproducen en la práctica. Comparto el código de resultados del proyecto mirage. Con la combinación Triton + FlashInfer, la latencia por token fue de alrededor de 19.2 ms, mientras que con MPK mejoró drásticamente a 7.7 ms bajo las mismas condiciones.
Hace tiempo participé en una pequeña competencia de CUDA. Era un algoritmo paralelo del área de imagen o visión, y para verme listo cacheé los resultados intermedios en memoria. Cuando vi los resultados del concurso, me sorprendió que otros hubieran enviado código mucho más rápido que el mío. Al revisar el motivo, vi que no cacheaban esos resultados intermedios y simplemente los recalculaban una y otra vez. El costo computacional era mucho menor que ir y volver a memoria. Supongo que este proyecto probablemente sea parecido. Al compilar a megakernel, desaparecen los límites entre capas, así que se comparte menos el resultado intermedio y aumenta la cantidad de cómputo, pero en conjunto hay una gran ganancia al reducir los viajes a memoria. Especialmente en redes convolucionales debe haber algún sweet spot, aunque no sé cómo maneja eso el megakernel.
Siguen apareciendo nuevas metáforas para los LLM. Se me ocurre si no podríamos pensar en los LLM como si fueran transistores. Me imagino que ahora estamos en una etapa parecida a la de las computadoras del tamaño de una habitación que solo hacían multiplicaciones con tarjetas perforadas. Es divertido imaginar qué pasaría si se pudieran ejecutar al mismo tiempo 1 millón de consultas a o3-pro.
Este proyecto viene de CMU (Carnegie Mellon). También se menciona el blog de Hazy Research de Stanford sobre megakernels, No Bubbles. Es impresionante ver lo activa que está la competencia en esta área. (Añadido) También hay un paper sobre el panorama más amplio del proyecto "mirage", aunque no trata el enfoque de megakernel: enlace al paper
El propio autor de la publicación responde directamente. Coincide en que la investigación con Stanford se está desarrollando en paralelo. La diferencia principal es que ellos se enfocan en un compilador para generación automatizada de megakernels.
Se menciona que ThunderKittens de Hazy Research también es una librería muy genial. Últimamente se está concentrando mucho esfuerzo en formalización, pipelining, divide y vencerás, maximización de eficiencia y desarrollo de compiladores/DSL dedicados para aprovechar al máximo los modelos recientes de GPU de NVIDIA.
Si se verifican las cifras de rendimiento de Qwen 8B, serían bastante impresionantes. Se siente más práctico que enfoques anteriores de megakernel. Este tipo de kernel que se mantiene uno por cada SM me recuerda a Larrabee. Me pregunto cómo habría sido el mundo actual si se hubiera seguido una ruta tradicional de proceso-hilo-SIMD, en vez de CUDA.
Idea sobre crear un LLM fijo en puro ASIC en vez de inferencia basada en software. ¿Ventaja en costos? ¿Posibilidad de ofrecer capas adicionales que el software pueda seguir manejando o afinando? Dado que en la práctica ya casi estamos en un nivel de “lo suficientemente bueno”, tal vez en los próximos 2 a 4 años alguien decida fijarlo en chips especializados para usarlo así. Me pregunto en qué punto empezarían realmente a destacar las ventajas que podría ofrecer hardware ultrasespecializado.