Resumen Técnico: Solapamiento de Activación a Nivel de Tile para la Inferencia Eficiente de LLM

Declaración del Problema

Los Modelos de Lenguaje Grandes (LLM) modernos, incluyendo Qwen-2.5, LLaMA, Mistral y Gemma, utilizan predominantemente la función de activación SwiGLU. Esta estructura de puerta (gated) de MLP computa $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$, lo que requiere dos multiplicaciones de matrices independientes (GeMMs) seguidas de una activación de puerta elemento a elemento.

En las implementaciones estándar, este proceso materializa dos tensores intermedios completos (Gate y Up) en la Memoria de Alto Ancho de Banda (HBM) entre las etapas de GeMM y de activación. El perfilado en GPUs NVIDIA H100 revela que esta materialización intermedia y la computación de la activación asociada consumen entre el 9% y el 37% del tiempo total de ejecución de la MLP. La sobrecarga es más severa para modelos más pequeños (por ejemplo, Qwen-2.5 0.5B), donde SwiGLU representa más del 30% del tiempo de la MLP. A medida que la densidad aritmética de los Tensor Cores aumenta mediante cuantización (FP8, INT4), el costo relativo de estas operaciones limitadas por memoria crece, creando un cuello de botella significativo.

La infraestructura de compilación existente, específicamente torch.compile de PyTorch, no logra abordar esto. No puede fusionar operaciones a través de dos GeMMs separados con matrices de pesos diferentes debido a limitaciones fundamentales en los pasos de fusión a nivel de grafo. Los experimentos muestran que torch.compile alcanza solo entre el 34% y el 94% del rendimiento de PyTorch eager para este patrón, y los indicios de fusión explícitos proporcionan una mejora insignificante.

Metodología

Los autores presentan dos núcleos SM90 basados en CUTLASS complementarios y diseñados a mano que fusionan SwiGLU en GeMM al nivel de tile, eliminando la materialización intermedia en HBM. Ambos núcleos aprovechan el Acelerador de Memoria de Tensor (TMA) y la programación especializada por warps disponible en la H100 (SM90).

1. Kernel-1: Solapamiento de Sincronización SwiGLU mediante Programación Pingpong

• Mecanismo: Este núcleo utiliza la programación especializada por warps Pingpong de CUTLASS, que particiona los warps en productores (carga de datos vía TMA) y consumidores (ejecución de MMA).
• Estrategia: Crea una cuadrícula de threadblocks $[M, N]$. Dentro de la fase del consumidor, después de que se computa el tile de Gate pero antes de que lleguen los datos del tile de Up, existe una breve ventana donde los warps consumidores están inactivos. El núcleo explota esta ventana para computar la activación Swish ($\text{SiLU}(Gate)$) en el acumulador de Gate residente en los registros.
• Solapamiento: La computación de Swish se solapa temporalmente con las cargas TMA de los tiles de peso de Up. La multiplicación elemento a elemento final ($\text{SiLU}(Gate) \odot Up$) y el almacenamiento ocurren en el epílogo.
• Optimización: Este diseño está optimizado para tamaños de lote grandes ($M \ge 2048$) donde la cuadrícula $[M, N]$ satura suficientemente los Procesadores Multihilo (SM) de la GPU.

2. Kernel-2: Intercalado de SwiGLU mediante un Árbol de Visitantes de Epílogo Personalizado

• Mecanismo: Este núcleo introduce un nodo de Árbol de Visitantes de Epílogo (EVT) personalizado llamado PairMulStore.
• Estrategia: Concatena las matrices de peso $W_1$ y $W_2$ en una única matriz fusionada $W_{fused} \in \mathbb{R}^{K \times 2N}$ con columnas intercaladas (columnas pares para Gate, impares para Up). El núcleo lanza un GeMM estándar sobre $[M, 2N, K]$, creando una cuadrícula de threadblocks $[M, 2N]$.
• Fusión: El nodo PairMulStore opera sobre el acumulador $[M, 2N]$ dentro de cada threadblock. Lee pares de columnas adyacentes $(D[m, 2n], D[m, 2n+1])$, aplica $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ y almacena el resultado.
• Solapamiento: Este diseño intercala la computación de SwiGLU con los almacenamientos de tiles. La aritmética del tile actual se solapa con los almacenamientos TMA asíncronos del tile anterior, proporcionando un solapamiento de cómputo-almacenamiento.
• Optimización: El aumento de $2\times$ en el conteo de threadblocks mejora significativamente la ocupación de los SM para tamaños de lote pequeños ($M \le 2048$), convirtiéndolo en el predeterminado preferido para la mayoría de las configuraciones.

Contribuciones Clave

1. Fusión a Nivel de Tile: La primera implementación que fusiona SwiGLU en GeMM al nivel de tile, eliminando 4 operaciones de HBM intermedias (lectura/escritura de Gate y Up) por cada paso de inferencia.
2. Estrategia de Doble Núcleo: Un enfoque complementario utilizando el Kernel-1 para lotes grandes (secuencia secuencial Gate $\to$ Swish $\to$ Up) y el Kernel-2 para lotes pequeños (pesos intercalados con EVT personalizado), abordando las restricciones de ocupación y programación de manera distinta.
3. Superioridad Numérica: Los núcleos fusionados logran cero discrepancias (0% de elementos con un error relativo >1%) en comparación con una referencia en FP32, mientras que la implementación base de cuBLAS exhibe entre un 4.5% y un 11% de discrepancias debido a la acumulación parcial agresiva en BF16 a través de los tiles.
4. Validación de Limitaciones del Compilador: Evidencia empírica que demuestra que torch.compile no puede replicar esta fusión, siendo de 3 a 7 veces más lento que los núcleos diseñados a mano.

Resultados

Evaluado en modelos Qwen-2.5 (0.5B a 72B) en NVIDIA H100:

• Aceleración (Speedup): Los núcleos logran hasta un 2.47× de aceleración sobre PyTorch eager (cuBLAS + SwiGLU separado). El Kernel-2 gana en 13 de 20 configuraciones evaluadas, mientras que el Kernel-1 gana en 7 de 20 (principalmente para modelos grandes con lotes altos).
• Cambio de Régimen: El análisis de Roofline confirma que la fusión desplaza las cargas de trabajo de limitadas por memoria a limitadas por cómputo. Para modelos pequeños, la intensidad aritmética aumenta hasta un 247%, moviendo el punto de operación a través de la cresta del roofline.
• Utilización: La utilización pico de BF16 alcanza el 79.5% en H100, casi duplicando la utilización de la base para modelos pequeños.
• Eficiencia: La eficiencia de fusión (aceleración real frente a ahorros teóricos de memoria) es cercana al 100% para modelos pequeños. Para el Kernel-2, la eficiencia supera el 100% en varias configuraciones debido a los beneficios del solapamiento de cómputo-almacenamiento.
• Escalabilidad: La aceleración disminuye con el tamaño del modelo (2.47× para 0.5B frente a ~1.0× para 72B en lotes altos) porque el cómputo de GeMM domina el tiempo total y la fracción relativa del tráfico de memoria intermedia disminuye.

Significado

El artículo argumenta que para la inferencia de LLM modernos, particularmente en modelos de despliegue en el borde (edge) y de tamaño móvil (0.5B–14B), el cuello de botella ha pasado de la computación pura al tráfico de memoria causado por la materialización de tensores intermedios. Los autores demuestran que el diseño a nivel de tile diseñado a mano y consciente del hardware es necesario para superar las limitaciones de la infraestructura de compilación actual. Al eliminar los accesos intermedios a HBM y solapar la computación de la activación con el movimiento de datos, estos núcleos reducen significativamente la latencia de inferencia y mejoran la utilización del cómputo, ofreciendo una optimización crítica para sistemas de servicio de producción donde la latencia y el rendimiento están restringidos.