LP-GEMM: Integrating Layout Propagation into GEMM Operations

Resumen Técnico: LP-GEMM

1. El Problema

En la computación científica y las cargas de trabajo modernas de Aprendizaje Automático (ML), como los modelos de Transformadores (LLM) y Perceptrones Multicapa (MLP), el tiempo de ejecución está dominado por secuencias de multiplicaciones de matrices generales (GEMM) dependientes.

Aunque las bibliotecas BLAS de última generación (como OpenBLAS, Intel MKL) optimizan agresivamente cada llamada individual a GEMM, operan bajo una restricción fundamental de la API BLAS: cada llamada debe ser independiente. Esto implica que:

1. La biblioteca debe empaquetar (pack) las matrices de entrada en un formato contiguo y optimizado para el micro-kernel.
2. Al finalizar la operación, debe desempaquetar (unpack) el resultado para restaurar el formato de memoria canónico estándar.

En secuencias de GEMMs (donde la salida de una operación es la entrada de la siguiente), este enfoque provoca empaquetado y desempaquetado redundantes en cada frontera de operación. Estos pasos de movimiento de datos consumen recursos computacionales valiosos sin aportar valor al cálculo matemático, limitando el rendimiento general de la carga de trabajo.

2. Metodología: LP-GEMM

El artículo presenta LP-GEMM (Layout Propagation GEMM), una descomposición del kernel GEMM diseñada para propagar el formato de datos (layout) entre operaciones secuenciales, eliminando la redundancia.

Concepto Central:
En lugar de tratar cada GEMM como una isla, LP-GEMM divide el proceso en tres tipos de kernels especializados que mantienen un formato de datos interno consistente:

1. Initial GEMM (Ini-GEMM): Realiza el empaquetado inicial de los datos de entrada y establece el "layout propagado" (un formato optimizado para el micro-kernel).
2. Intermediate GEMM (Mid-GEMM): Recibe los datos ya empaquetados desde la operación anterior. Omite el paso de empaquetado y mantiene el mismo layout para la siguiente operación.
3. Ending GEMM (End-GEMM): Realiza el último cálculo y desempaqueta el resultado final para devolverlo al formato de memoria estándar requerido por el usuario o la siguiente etapa no-GEMM.

Implementación Técnica:

• Propagación de Layout: El enfoque alinea tres aspectos que en BLAS tradicional están desacoplados: (1) el layout empaquetado para lectura, (2) el orden de iteración durante el cálculo, y (3) el layout final almacenado. En LP-GEMM, estos tres son idénticos, permitiendo que la salida de un kernel sea la entrada directa del siguiente sin reordenamiento.
• Micro-kernels: Se utilizan dos variantes de micro-kernels:
  • Propagate-Layout µkernel: Mantiene el layout propagado.
  • Default µkernel: Restablece el layout estándar (usado en el kernel final).
• Adaptabilidad: Se han implementado funciones para manejar accesos a memoria con saltos (strided loads/stores) y se han adaptado operaciones intermedias comunes en ML (como Softmax, RoPE y RMSNorm) para funcionar eficientemente sobre el layout propagado, minimizando la sobrecarga de traducción de direcciones.

3. Contribuciones Clave

• Descomposición de Kernels: Introducción de los kernels ini-GEMM, mid-GEMM y end-GEMM derivados de OpenBLAS, que permiten la propagación de layouts.
• Implementación de LLaMA 3.2: Desarrollo de una versión autónoma en C++ del camino de inferencia de LLaMA 3.2 (1B parámetros) utilizando exclusivamente llamadas GEMM de nivel BLAS, validando la viabilidad práctica del enfoque.
• Adaptación de Operadores ML: Implementación de operaciones críticas (Softmax, RoPE, RMSNorm) optimizadas para el layout propagado en arquitecturas AVX-512.
• Validación Multi-arquitectura: Evaluación exhaustiva en arquitecturas x86 (AVX-512) y RISC-V (RVV 1.0).

4. Resultados Experimentales

Los experimentos se realizaron comparando LP-GEMM contra OpenBLAS, BLIS, Intel MKL, oneDNN y FlashGEMM.

• GEMM Aislado (Micro-benchmarks):
  • En x86, los kernels Mid y End lograron una aceleración mediana de 1.5x sobre OpenBLAS, con cuartiles superiores a 2x. Esto demuestra que el empaquetado es a menudo el costo dominante.
  • En RISC-V, las ganancias fueron aún más significativas, alcanzando hasta 5x en cargas de trabajo tipo Atención, debido a la ineficiencia del kernel de desempaquetado de OpenBLAS en esta arquitectura.
• Carga de Trabajo Real (Capa de Atención de LLaMA 3.2):
  • En x86, se observaron aceleraciones de hasta 2.25x en secuencias de GEMMs. El rendimiento escala inversamente con el tamaño de la matriz (mayor beneficio en matrices más pequeñas/medias donde el costo de empaquetado es proporcionalmente mayor).
  • En RISC-V, la aceleración aumentó linealmente con el tamaño de la secuencia de tokens, superando significativamente a las bibliotecas de referencia.
• Secuencias de GEMMs Consecutivos:
  • En benchmarks de tres GEMMs consecutivos (extraídos de capas convolucionales), LP-GEMM superó consistentemente a OpenBLAS y a FlashGEMM (una solución de fusión agresiva), demostrando que la propagación de layout es una estrategia robusta incluso sin fusión completa de kernels.

5. Significado y Conclusión

LP-GEMM representa un cambio de paradigma en la optimización de álgebra lineal:

• Eficiencia sin Tuning de Bajo Nivel: Logra mejoras de rendimiento significativas no mediante un ajuste micro-arquitectural más profundo del kernel de multiplicación, sino eliminando la sobrecarga de movimiento de datos (empaquetado/desempaquetado).
• Portabilidad: Al basarse en la estructura de OpenBLAS, el enfoque es portable a diferentes arquitecturas (x86, RISC-V, ARM, etc.) sin necesidad de reescribir los micro-kernels desde cero.
• Viabilidad Práctica: La implementación de LLaMA 3.2 demuestra que este enfoque es aplicable en sistemas de inferencia reales, ofreciendo una alternativa viable a las bibliotecas propietarias y altamente optimizadas de los fabricantes.

El trabajo concluye que la optimización de cargas de trabajo de ML debe evolucionar desde la optimización de llamadas GEMM aisladas hacia la optimización de flujos de trabajo end-to-end, donde la propagación de layouts es una técnica fundamental para desbloquear el verdadero potencial del hardware moderno.