Mejora del autoajuste de GEMM con nvMatmulHeuristics en CUTLASS 4.2

Visión general

La selección de kernels GEMM en GPUs de NVIDIA modernas es un problema de optimización con múltiples parámetros. El rendimiento de un kernel depende de un conjunto amplio de meta-parámetros de compilación y ejecución, incluyendo tamaños de CTA, warp y tiles de instrucción; planificaciones de kernels; estrategias de rasterización; dimensiones de cluster; y factores de división (split-k). El enfoque tradicional genera miles de configuraciones, las compila y ejecuta una búsqueda exhaustiva para identificar la más rápida. Este flujo puede tomar horas y es especialmente problemático para bibliotecas fuera de línea como CUTLASS, y para entornos JIT como Torch Inductor u OpenAI Triton, donde la compilación rápida es crucial. NVIDIA introduce nvMatmulHeuristics, un módulo de optimización de meta-parámetros de kernels que ofrece heurísticas rápidas para predecir un pequeño conjunto de configuraciones GEMM prometedoras. El módulo analiza los parámetros de un problema GEMM y las capacidades de hardware objetivo para predecir configuraciones que entregan el máximo rendimiento. La integración con CUTLASS 4.2 transforma el flujo de generación de kernels y tuning de un brute-force hacia una ruta enfocada y eficiente. Para un problema GEMM dado, nvMatmulHeuristics genera un conjunto conciso de configuraciones candidatas para probar. El rasgo se posiciona como una parte central de las heurísticas de cuBLAS y está disponible en acceso anticipado para uso general, con integración en la biblioteca CUTLASS. Un ejemplo presenta un GEMM FP16 único (A transpuesto o no, B en formato columna mayor y C/D según layouts). El flujo implica proporcionar una lista de problemas GEMM en JSON y compilar CUTLASS con banderas de CMake específicas para emitir un conjunto de configuraciones por problema. En la práctica, este enfoque reduce drásticamente el tiempo total de ajuste. En lugar de compilar y perfilar exhaustivamente miles de kernels, se genera un conjunto reducido y dirigido y se ejecuta solo ese conjunto. Esto es especialmente valioso en entornos como PyTorch, donde los tiempos de compilación JIT son críticos, y para bibliotecas que desean un rendimiento cercano al óptimo sin largos ciclos de tuning.

Características clave

Predice un conjunto pequeño y top-N de configuraciones GEMM por problema
Reduce el tiempo total de tuning concentrando compilación y perfilado en un conjunto de candidatos
Integrado con CUTLASS 4.2 y heurísticas de cuBLAS; disponible en acceso anticipado
Acepta una entrada JSON que describe problemas GEMM; fácil de integrar en flujos existentes
Genera una lista CSV de pruebas a partir de las configuraciones emitidas para profiling automático
Compatible con cutlass_profiler para ejecutar las configuraciones generadas
Permite tamaños de cluster estáticos a la compilación, reduciendo la variabilidad en ejecución
Demuestra rendimiento cercano al máximo con mucho menos esfuerzo; por ejemplo, 16 configuraciones alcanzan 96% del pico con un tiempo significativamente menor
Ayuda a acelerar flujos DL, compiladores y bibliotecas de kernels para lograr alto rendimiento de manera eficiente

Tabla: comparación de enfoques de tuning

| Enfoque | Cómo funciona | Beneficios típicos | Observaciones |---|---|---|---| | Búsqueda exhaustiva | Genera y compila miles de kernels; tuning exhaustivo | El mejor kernel posible; rendimiento teórico máximo | Duración de construcción y tuning muy alta (ejemplos > 700 minutos) |nvMatmulHeuristics (esta función) | Predice top-N kernels por problema; compila un conjunto pequeño y realiza profiling | Ahorros de tiempo sustanciales; rendimiento cercano al pico | Requiere lista de problemas en JSON e integración de compilación |Precompilados/cluster estático | Usa tamaños de cluster fijos de kernels precompilados | Configuración rápida; profiling en ejecución reducido | Menos flexibilidad; desempeño depende de elecciones de compilación |

Casos de uso comunes

Integración en frameworks DL: acelerar rutas JIT (p. ej., PyTorch) reduciendo el tiempo de tuning, manteniendo alto rendimiento para cargas GEMM.
Bibliotecas fuera de línea: permitir conjuntos de kernels bien ajustados con mínimo esfuerzo de tuning durante la implementación.
Cargas de modelos con grandes GEMMs (p. ej., Llama, DeepSeek) donde el tuning exhaustivo resulta prohibitivo.
Escenarios que requieren perfiles reproducibles; el enfoque favorece un conjunto reducido de candidatos para builds reproducibles.

Configuración e instalación (comandos exactos)

El flujo se basa en preparar una lista de problemas GEMM en JSON, construir CUTLASS con nvMatmulHeuristics activado y luego ejecutar el profiling con el CSV generado. Las opciones explícitas de CMake que dirigen el flujo heurístico son:

cmake .. \
-DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE= \
-DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM=N

Tras configurar, una construcción estándar emitirá un CSV de pruebas durante el proceso. Puede usar cutlass_profiler para ejecutar las configuraciones generadas a partir del CSV; para obtener resultados de profiling consistentes, ejecute con relojes bloqueados. Notas:

Construya CUTLASS como de costumbre, pero con las dos banderas CUTLASS_LIBRARY_HEURISTICS indicadas arriba.
El formato exacto de los problemas JSON no está especificado aquí; el flujo se refiere a una lista de problemas GEMM expresados en JSON.

Inicio rápido (ejemplo mínimo ejecutable)

Prepare una lista de problemas GEMM en JSON (el ejemplo en el texto menciona un GEMM FP16 único, A transpuesto o no, y disposiciones de B y C/D).
Construya CUTLASS con heurísticas activadas usando las comandos de la sección anterior.
El sistema de compilación emitirá un CSV describiendo las configuraciones generadas para cada problema GEMM.
Ejecute el flujo de profiling con cutlass_profiler sobre el CSV generado para ejecutar las configuraciones. Para resultados consistentes, use relojes bloqueados durante el profiling.
Compare el rendimiento con una base exhaustiva, si está disponible; en los casos publicados, 16 configuraciones alcanzan el 96% del pico con mucho menos tiempo, y más candidatos se acercan al baseline exhaustivo manteniendo beneficios temporales.

Pros y contras

Pros
Reducción enorme del tiempo de build+tuning frente a una búsqueda exhaustiva
Rendimiento cercano al pico con un conjunto pequeño de candidatos
Mejor experiencia para pipelines JIT y despliegues en dispositivos
Opciones de cluster estático pueden mejorar la previsibilidad del compile-time
Contras
Requiere formato JSON para la lista de problemas e integración correspondiente
Recurso en acceso anticipado; las API pueden evolucionar
Dependiente de capacidades de hardware y de la precisión de la caracterización del problema

Alternativas (comparaciones breves)

Búsqueda exhaustiva: garantiza el mejor kernel, pero implica un gran costo de tiempo y recursos (los ejemplos citados muestran >700 minutos).
Con cluster estático con heurísticas: ofrece buen rendimiento, pero depende de las elecciones de compilación y puede reducir la flexibilidad en tiempo de ejecución.
Kernels precompilados/dinámicos tipo Blackwell: permiten adaptación en tiempo real, pero pueden requerir tuning en ejecución; nvMatmulHeuristics intenta trasladar más decisiones al tiempo de compilación.

Precio o licencia

nvMatmulHeuristics se describe como parte central de las heurísticas de cuBLAS y está disponible en acceso anticipado para uso general, con integración en CUTLASS. No se proporcionan detalles de precio en la fuente.