DeepGEMM - Biblioteca eficiente de multiplicación de matrices FP8
Una biblioteca clara y eficiente de multiplicación de matrices FP8 con escalado de grano fino, optimizada para la arquitectura NVIDIA Hopper
Diseño de arquitectura DeepGEMM
DeepGEMM presenta una arquitectura cuidadosamente diseñada optimizada para los Tensor Cores de NVIDIA Hopper, permitiendo una multiplicación de matrices FP8 eficiente.
Diseño completamente JIT
DeepGEMM emplea un diseño Just-In-Time (JIT) completo sin compilación en tiempo de instalación. Todos los kernels se compilan en tiempo de ejecución utilizando un módulo JIT ligero, tratando las formas GEMM, tamaños de bloque y etapas de pipeline como constantes en tiempo de compilación, ahorrando registros y permitiendo más optimizaciones del compilador.
Características Hopper TMA
Aprovecha completamente el Acelerador de Memoria Tensorial (TMA) de la arquitectura Hopper para un movimiento de datos más rápido y asíncrono, incluyendo cargas TMA, almacenamientos, difusiones y prefetch de descriptores, mejorando significativamente el rendimiento.
Técnicas de optimización principales
DeepGEMM emplea varias técnicas de optimización avanzadas que van más allá de las bibliotecas GEMM tradicionales para lograr un rendimiento excepcional.
Especialización de hilos persistentes
Siguiendo el diseño CUTLASS, superpone el movimiento de datos, las instrucciones MMA del tensor core y la promoción del núcleo CUDA para mejorar la eficiencia computacional.
Tamaños de bloque no alineados
Soporta tamaños de bloque no alineados (como 112) para utilizar mejor los SMs. Por ejemplo, con M=256, N=7168, usar BLOCK_M=128, BLOCK_N=112 permite que trabajen 128 SMs en lugar de 112 con tamaños de bloque alineados.
Entrelazado FFMA SASS
Modifica las instrucciones FFMA en el binario compilado invirtiendo los bits de yield y reuse, creando más oportunidades para superponer instrucciones MMA con instrucciones FFMA de promoción, produciendo un rendimiento 10%+ mejor en algunos casos.
Planificador de bloques unificado optimizado
Un planificador para todos los kernels no agrupados y agrupados, empleando rasterización para mejorar la reutilización de caché L2, mejorando el rendimiento general.
Código explicado
La función kernel principal de DeepGEMM tiene solo unas 300 líneas de código, con un diseño limpio que facilita el aprendizaje sobre multiplicación de matrices FP8 en Hopper y técnicas de optimización.
// Simplified version of the core GEMM kernel
template <
int BLOCK_M,
int BLOCK_N,
int BLOCK_K,
int CLUSTER_M,
int CLUSTER_N,
int STAGES,
int WARPGROUP_SIZE,
bool ENABLE_INTERLEAVE
>
__global__ void gemm_fp8_fp8_bf16_nt_kernel(
const void* __restrict__ lhs,
const void* __restrict__ lhs_scale,
const void* __restrict__ rhs,
void* __restrict__ output,
const void* __restrict__ rhs_scale,
int m, int n, int k,
int lhs_stride, int rhs_stride, int output_stride) {
// Using TMA to load LHS, LHS scale factors, and RHS
// Using CUDA cores for two-level accumulation (promotion)
// Using TMA to store output
// Persistent thread specialization implementation
// Unaligned block size support
// FFMA SASS interleaving optimization
}Características de la estructura del código
- Diseño basado en plantillas que soporta diferentes tamaños de bloque y configuraciones de optimización
- Utilización completa de las características TMA de la arquitectura Hopper para un movimiento de datos eficiente
- Acumulación de dos niveles para resolver el problema de acumulación imprecisa de los tensor cores FP8
- Implementación de especialización de hilos persistentes para computación optimizada y superposición de movimiento de datos
Diseño de interfaz
- GEMM denso normal: Usando la función deep_gemm.gemm_fp8_fp8_bf16_nt
- GEMM agrupado (Diseño contiguo): Usando la función m_grouped_gemm_fp8_fp8_bf16_nt_contiguous, adecuada para modelos MoE donde los expertos comparten la misma forma
- GEMM agrupado (Diseño enmascarado): Usando la función m_grouped_gemm_fp8_fp8_bf16_nt_masked, utilizada en la fase de decodificación de inferencia con gráficos CUDA habilitados donde la CPU no conoce el número de tokens que recibe cada experto
Rendimiento
El rendimiento de DeepGEMM es comparable o mejor que las bibliotecas ajustadas por expertos en varias formas de matrices, con hasta 2.7x de aceleración en algunos casos.
Aspectos destacados del rendimiento
- GEMM normal: Hasta 2.7x de aceleración en ciertas formas
- GEMM agrupado con diseño contiguo: Hasta 1.2x de aceleración
- GEMM agrupado con diseño enmascarado: También hasta 1.2x de aceleración
- Rendimiento computacional de hasta 1358 TFLOPS probado en H800 SXM5
- Alta utilización del ancho de banda de memoria, alcanzando hasta 2668 GB/s
Características principales de DeepGEMM
Rendimiento excepcional
DeepGEMM ofrece hasta 2.7x de aceleración en comparación con bibliotecas ajustadas por expertos en la arquitectura NVIDIA Hopper, logrando más de 1350 TFLOPS
Optimizaciones avanzadas
DeepGEMM emplea especialización de hilos persistentes, tamaños de bloque no alineados y entrelazado FFMA SASS para máxima eficiencia computacional
Integración flexible
DeepGEMM se integra perfectamente con frameworks de aprendizaje profundo y bibliotecas de computación científica a través de interfaces claras y bien documentadas
Escalado de grano fino
DeepGEMM escala eficientemente desde matrices pequeñas hasta cálculos a gran escala con planificación de bloques optimizada y técnicas de rasterización
Soporte nativo FP8
DeepGEMM proporciona soporte de primera clase para precisión FP8, optimizado específicamente para los Tensor Cores de NVIDIA Hopper
Alto ancho de banda de memoria
DeepGEMM logra una utilización excepcional del ancho de banda de memoria de hasta 2668 GB/s, maximizando las capacidades del hardware
Diseño de código limpio
DeepGEMM presenta un kernel central conciso de 300 líneas con patrones de diseño claros, facilitando su comprensión y extensión
GEMM agrupado optimizado
DeepGEMM sobresale en operaciones GEMM agrupadas tanto de diseño contiguo como enmascarado, perfecto para modelos de Mezcla de Expertos (MoE)
Preguntas frecuentes
- ¿Qué es DeepGEMM y qué lo hace especial?
- DeepGEMM es una biblioteca eficiente de multiplicación de matrices FP8 optimizada para la arquitectura NVIDIA Hopper. Lo que lo hace especial es su diseño limpio (solo ~300 líneas de código central), rendimiento excepcional (hasta 2.7x de aceleración sobre bibliotecas ajustadas por expertos) y técnicas de optimización avanzadas como especialización de hilos persistentes y entrelazado FFMA SASS.
- ¿Para qué hardware está optimizado DeepGEMM?
- DeepGEMM está específicamente optimizado para la arquitectura NVIDIA Hopper, aprovechando al máximo las características del Acelerador de Memoria Tensorial (TMA) de Hopper y los Tensor Cores para multiplicación de matrices FP8. Ha sido probado en GPUs H800 SXM5, logrando hasta 1358 TFLOPS de rendimiento computacional.
- ¿Cómo logra DeepGEMM sus ganancias de rendimiento?
- DeepGEMM logra su rendimiento a través de varias técnicas avanzadas: diseño completamente JIT con constantes en tiempo de compilación, tamaños de bloque no alineados para mejor utilización de SM, entrelazado FFMA SASS, acumulación de dos niveles para resolver problemas de precisión FP8, y un planificador de bloques optimizado unificado con rasterización para mejorar la reutilización de caché L2.
- ¿Qué tipos de operaciones GEMM soporta DeepGEMM?
- DeepGEMM soporta tres tipos principales de operaciones GEMM: GEMM denso normal (usando la función deep_gemm.gemm_fp8_fp8_bf16_nt), GEMM agrupado con diseño contiguo, y GEMM agrupado con diseño enmascarado. Los dos últimos son particularmente útiles para modelos de Mezcla de Expertos (MoE).
- ¿Es DeepGEMM adecuado para modelos de Mezcla de Expertos (MoE)?
- Sí, DeepGEMM es altamente adecuado para modelos MoE. Proporciona operaciones GEMM agrupadas especializadas con diseños tanto contiguos como enmascarados, que están optimizadas para los patrones de cálculo específicos en arquitecturas MoE. Estas operaciones pueden ofrecer hasta 1.2x de aceleración en comparación con otras bibliotecas.
- ¿Qué precisión soporta DeepGEMM?
- DeepGEMM se centra principalmente en la precisión FP8, que es cada vez más importante para la inferencia eficiente de modelos de IA. Implementa un enfoque de acumulación de dos niveles para resolver los problemas de acumulación imprecisa que pueden ocurrir con los tensor cores FP8, asegurando tanto rendimiento como precisión.
- ¿Cómo se compara la utilización del ancho de banda de memoria de DeepGEMM con otras bibliotecas?
- DeepGEMM logra una utilización excepcional del ancho de banda de memoria, alcanzando hasta 2668 GB/s. Esto es posible gracias a su uso eficiente de las características TMA de Hopper para un movimiento de datos más rápido y asíncrono, incluyendo cargas TMA, almacenamientos, difusiones y prefetch de descriptores.
- ¿Puedo aprender de la implementación de DeepGEMM para optimizar mi propio código CUDA?
- ¡Absolutamente! El kernel central de DeepGEMM tiene solo unas 300 líneas de código con un diseño limpio, lo que lo convierte en un excelente recurso de aprendizaje. Demuestra técnicas avanzadas de optimización CUDA, uso eficiente de las características de la arquitectura Hopper, y enfoques inteligentes para la multiplicación de matrices que pueden aplicarse a otras tareas de computación de alto rendimiento.