Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

Resumen Técnico: Aceleración de la Inferencia de Modelos de Lenguaje Extensos en GPU con Matrices de Pesos con Esparcidad Moderadamente No Estructurada

Planteamiento del Problema
El despliegue de Modelos de Lenguaje Extensos (LLMs) enfrenta un cuello de botella crítico en los costos de inferencia. Aunque las técnicas de poda de modelos (p. ej., Wanda, RIA) pueden reducir la carga computacional al introducir esparcidad en las matrices de pesos, están limitadas por los requisitos de calidad del modelo. A diferencia de los modelos de pequeña escala, los LLMs son altamente sensibles a niveles altos de esparcidad; mantener el rendimiento suele limitar la poda a una esparcidad no estructurada moderada (alrededor del 50%).

En este nivel específico de esparcidad, los kernels de GPU existentes para la Multiplicación de Matrices Dispersas (SpMM) no logran superar a las librerías de multiplicación de matrices densas como cuBLAS. Esta brecha de rendimiento surge de tres desafíos fundamentales:

1. Incompatibilidad con los Tensor Cores: Los GPUs modernos ofrecen Tensor Cores dispersos de alta eficiencia, pero estos soportan estrictamente la esparcidad estructurada 2:4. No pueden acelerar directamente los patrones de poda irregulares o no estructurados.
2. Sobrecarga de Metadatos: Los formatos de almacenamiento tradicionales (p. ej., CSR) requieren almacenar metadatos posicionales completos para cada elemento no nulo. Con una esparcidad del 50%, estos metadatos pueden ocupar tanto espacio como los propios valores, anulando los beneficios del ancho de banda de memoria.
3. Cuellos de Botella de Decodificación en Chip: Las técnicas de compresión que reducen los metadatos (p. ej., codificación por mapa de bits) trasladan la carga de trabajo a la decodificación en chip. En los GPUs modernos con Memoria de Alto Ancho de Banda (HBM), el acceso a la memoria global no puede solaparse completamente con el trabajo de decodificación realizado por los más lentos núcleos CUDA, lo que crea un estancamiento en el rendimiento.

Metodología
Los autores proponen una solución holística que comprende un nuevo formato de almacenamiento de tres capas y un kernel de SpMM co-optimizado diseñado para la esparcidad no estructurada moderada.

1. Formato de Almacenamiento de Tres Capas
El formato propuesto descompone la matriz dispersa en tres capas distintas para equilibrar la compatibilidad con el hardware y la eficiencia de almacenamiento:

• Capa Sparse-TC: Esta capa extrae patrones de esparcidad estructurada 2:4 de la distribución irregular. Retiene dos valores no nulos de cada cuatro pesos consecutivos. Debido a que el patrón es fijo, requiere solo 2 bits por valor para los metadatos, asegurando la compatibilidad directa con los sparse tensor cores sin sobrecarga de decodificación.
• Capa de Llenado de Ranuras (Slot-Filling): Los elementos no nulos excedentes que no encajan en el patrón 2:4 se redistribuyen en los espacios vacíos dejados por la capa Sparse-TC. Para minimizar los metadatos de estos ~17.7% de los elementos, los autores introducen la codificación de Distancia Diferencial Paralela (PDD). En lugar de almacenar índices absolutos, la PDD codifica las distancias relativas entre los elementos redistribuidos. La mayoría de las distancias son lo suficientemente pequeñas como para codificarse en 4 bits. Para brechas más grandes, se insertan puntos ficticios para romper la extensión, asegurando que ninguna distancia única exceda el límite de codificación.
• Capa Residual: Los elementos no nulos restantes (típicamente <1% del total) que no pueden ser manejados eficientemente por las dos primeras capas se almacenan en un formato de Fila Dispersa Comprimida (CSR) tradicional. Debido al volumen mínimo de datos, la sobrecarga de metadatos aquí es insignificante.

2. Kernel Co-Optimizado
Basándose en este formato, los autores diseñan un kernel de GPU que utiliza conjuntamente sparse tensor cores y núcleos CUDA dentro de un modelo de ejecución en pipeline.

• Decodificación Paralela: Los núcleos CUDA se encargan de decodificar los metadatos PDD de la capa de Llenado de Ranuras en paralelo. Dado que esta capa representa una pequeña fracción del total de los datos, la sobrecarga de decodificación se gestiona de manera efectiva.
• Aceleración Estructurada: Los sparse tensor cores procesan la capa Sparse-TC utilizando instrucciones wgmma (arquitectura Hopper), aprovechando la estructura 2:4 para obtener un alto rendimiento.
• Manejo de Residuales: Los núcleos CUDA también procesan la capa Residual dispersa, escribiendo los resultados en un búfer separado en chip para ser reducidos con los resultados de los tensor cores al final del pipeline.
• Solapamiento de Pipeline: El kernel implementa un pipeline de grano fino que solapa tres operaciones críticas:
  • La carga de matrices de entrada densas (cacheadas en L2) y teselas de pesos dispersos (desde la Memoria Global).
  • La decodificación de metadatos PDD en los núcleos CUDA.
  • La ejecución de la multiplicación de matrices en los sparse tensor cores.
    Esta estrategia de "cargar-como-disperso, computar-como-disperso" asegura que la computación en chip se solape con el acceso a la memoria global, maximizando la utilización de la HBM.

Contribuciones Clave

• Nuevo Formato de Almacenamiento: Un formato de tres capas adaptado para la esparcidad no estructurada moderada que permite la compatibilidad con los sparse tensor cores, reduce drásticamente la sobrecarga de metadatos mediante la codificación PDD y minimiza los costos de decodificación en chip.
• Kernel Co-Optimizado: Un kernel de SpMM que coordina eficientemente los sparse tensor cores y los núcleos CUDA, utilizando el pipelining para solapar el acceso a la memoria con la computación.
• Primero en Superar a la Densidad: La implementación es la primera en demostrar que la SpMM con esparcidad moderadamente no estructurada puede superar la multiplicación de matrices densa en GPUs modernos equipados con HBM.

Resultados
Las evaluaciones se realizaron en GPUs NVIDIA H100 SXM5 utilizando los modelos OPT-30B y OPT-66B con niveles de esparcidad del 50%, 60% y 70%.

• Rendimiento a Nivel de Kernel: Con un 50% de esparcidad, el método propuesto logra una aceleración de hasta 1.64× sobre SpInfer (Best Paper de EuroSys'25) y supera significativamente a otros baselines como SparTA y FlashLLM. También supera el rendimiento de cuBLAS denso.
• Inferencia de Extremo a Extremo: El método ofrece una aceleración de extremo a extremo de hasta 1.41× respecto a FlashLLM (VLDB'24) y un throughput hasta un 41.4% mayor comparado con FlashLLM en OPT-30B con un tamaño de lote (batch size) de 16.
• Eficiencia de Memoria: El enfoque reduce el uso de memoria global pico en aproximadamente un 21.4% en comparación con los baselines densos al 50% de esparcidad, aunque incurre en un uso de memoria ligeramente superior al de SpInfer debido a la falta de codificación por mapa de bits para los metadatos.
• Escalabilidad: El método mantiene sus ventajas a través de los niveles de esparcidad del 50% al 70%.

Significancia y Reivindicaciones
El artículo afirma abordar una brecha crítica en la optimización de la inferencia de LLMs: la incapacidad de los kernels dispersos existentes para aprovechar eficazmente la esparcidad no estructurada moderada. Al introducir un formato de almacenamiento que cierra la brecha entre los patrones de poda irregulares y la aceleración de hardware estructurada, los autores demuestran que la poda puede generar aceleraciones reales de inferencia sin sacrificar la calidad del modelo.

El trabajo destaca que, si bien la esparcidad alta (>70%) degrada la calidad del modelo, la esparcidad moderada (~50%) es el "punto ideal" (sweet spot) para los LLMs. La solución propuesta desbloquea el potencial de rendimiento de este punto ideal, que anteriormente era inaccesible debido a las ineficiencias de los kernels dispersos existentes. Los autores señalan que su optimización está dirigida específicamente a la fase de decodificación de la inferencia de LLM (donde los tamaños de lote son pequeños y están limitados por memoria), lo que la hace altamente adecuada para el despliegue escalable en entornos de producción donde las fases de prefill y decode suelen estar desacopladas. El código fuente se pone a disposición para facilitar la investigación y adopción posteriores.