Mejorando la Eficiencia del Auto-Tuning de Kernels GEMM en GPUs NVIDIA con Heurísticas y CUTLASS 4.2

TL;DR

• nvMatmulHeuristics ofrece heurísticas rápidas para GEMM y predice un conjunto pequeño de configuraciones de kernel con alto potencial.
• Integrado a CUTLASS 4.2, reduce drásticamente el tiempo de tuning de kernels al pasar de una búsqueda exhaustiva a una evaluación dirigida de candidatos.
• Resultados iniciales: al usar 16 configuraciones, se alcanza aproximadamente el 96% del rendimiento máximo en unos 150 minutos, frente a más de 700 minutos para una búsqueda exhaustiva (Llama 3 405B en H100 SXM).
• Para DeepSeek-R1 671B en un GPU B200, 8 configuraciones alcanzan el 99% del rendimiento de la búsqueda exhaustiva con más de 5x de aceleración en el tiempo de construcción y auto-tuning. Esto permite flujos tipo JIT en marcos de DL.
• El enfoque favorece tamaños de cluster estáticos conocidos en tiempo de compilación, lo que puede entregar kernels más eficientes que con clustering dinámico en ciertos kernels precompilados. Estos resultados muestran cómo una heurística bien diseñada puede acercarse al rendimiento máximo de kernel mientras reduce drásticamente los costos de tuning. NVIDIA Developer Blog

Contexto y antecedentes

Elegir el mejor kernel GEMM para un problema y hardware específicos es un desafío de optimización. El rendimiento del GEMM depende de una amplia gama de metaparámetros de compilación y de tiempo de ejecución, como CTA, tamaños de tile por warp y a nivel de instrucción, agendas de kernel, estrategias de rasterización, dimensiones de cluster y factores de división. Tradicionalmente, encontrar el kernel óptimo implicaba generar miles de configuraciones, compilarlas y ejecutarlas en un tuning exhaustivo. Este flujo puede tomar horas y se convierte en un cuello de botella para bibliotecas offline como CUTLASS, así como para bibliotecas JIT como Torch Inductor u OpenAI Triton, donde la compilación rápida es clave. El artículo de NVIDIA presenta nvMatmulHeuristics, un módulo de optimización de meta-parametros para kernels GPU, diseñado para entregar un conjunto breve y de alto potencial de configuraciones de kernel para un GEMM y hardware dada. El módulo analiza los parámetros de la operación y las capacidades del hardware para prever un conjunto conciso de configuraciones cercanas al rendimiento máximo. La integración en CUTLASS 4.2 busca transformar el flujo de generación de kernels y tuning en un proceso más rápido y predecible. Esta característica forma parte de los heurísticos de cuBLAS y está disponible en acceso temprano para su uso general, con una integración en CUTLASS. NVIDIA Developer Blog El flujo pasa de una búsqueda por fuerza bruta a una exploración guiada. Los usuarios deben preparar una lista de problemas GEMM en formato JSON, construir CUTLASS con heurísticas habilitadas y especificar -DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE= y -DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM=N, donde N es el número de configuraciones que nvMatmulHeuristics emitirá por cada GEMM. La etapa de compilación genera una lista CSV de pruebas que enumera todos los casos que deben ejecutarse para realizar el tuning sobre las configuraciones emitidas. Esta CSV puede utilizarse con código de benchmarking personalizado o con cutlass_profiler para ejecutar las configuraciones de forma automática. Para resultados de profiling consistentes, se recomienda ejecutar con clocks bloqueados. La característica forma parte del ecosistema CUTLASS y es compatible con GEMMs FP16 de entrada/salida y FP32 de cómputo (HSH). NVIDIA Developer Blog

Un ejemplo FP16 y rendimiento práctico

Para un escenario FP16 de entrada/salida y cómputo FP32 (HSH), nvMatmulHeuristics puede emitir ocho configuraciones. Estos ocho kernels pueden compilarse y evaluarse, y seleccionar el mejor kernel entre ellos permite desplegar sin el costo de una búsqueda exhaustiva, manteniendo un rendimiento cercano al óptimo. En los resultados publicados, ocho configuraciones alcanzaron el 104% del baseline al probar GEMMs con tamaños de cluster estáticos, lo que demuestra que enfocar el esfuerzo de compilación y de perfil en un conjunto reducido puede lograr rendimiento de pico a menor costo que una búsqueda exhaustiva. NVIDIA Developer Blog

Resultados y conclusiones empíricas

Dos cargas de trabajo representativas ilustran el impacto de nvMatmulHeuristics. Primero, una carga de entrenamiento Llama 3 405B en una GPU NVIDIA H100 SXM mostró que la búsqueda exhaustiva puede tomar más de 700 minutos para identificar el kernel óptimo, mientras que una selección de 16 configuraciones alcanza 96% del rendimiento máximo en aproximadamente 150 minutos. Esto evidencia una reducción drástica del tiempo total de tuning con una caída menor respecto al rendimiento pico. En una segunda carga de DeepSeek-R1 671B en una GPU NVIDIA B200, se observó que usar un conjunto pequeño de candidatos puede lograr 99% del rendimiento de la búsqueda exhaustiva con más de 5x de aceleración en tiempo de construcción y auto-tuning. En este caso, baseline utilizaba tamaños de cluster dinámicos, mientras que las configuraciones propuestas por nvMatmulHeuristics se construyen con tamaños de cluster estáticos conocidos en tiempo de compilación, resultando en un rendimiento del 104% respecto al baseline para estos GEMMs. Estos resultados demuestran cómo heurísticas bien definidas pueden acercar el rendimiento a la cumbre para cargas de trabajo, reduciendo costos de tuning y facilitando la integración en frameworks, compiladores y bibliotecas de kernels. NVIDIA Developer Blog

Implicaciones para desarrolladores y empresas

Al enfocar el tuning en un conjunto reducido de kernels de alto potencial, nvMatmulHeuristics reduce considerablemente el tiempo requerido para desplegar GEMMs de alto rendimiento en entornos de producción. Esto facilita la adopción de kernels optimizados en stacks JIT y bibliotecas offline, acelerando los ciclos de desarrollo para frameworks de DL, compiladores y bibliotecas de kernels. La integración en CUTLASS 4.2 simplifica la adopción para desarrolladores que ya utilizan CUTLASS y heurísticas cuBLAS como parte de su stack. NVIDIA Developer Blog

Detalles técnicos o Implementación

nvMatmulHeuristics opera sobre dos capas: un módulo de heurísticas predictivas que analiza los parámetros del GEMM y las capacidades del hardware, y una integración con CUTLASS que utiliza las configuraciones previstas para construir y evaluar un conjunto compacto de kernels. Pasos clave:

• Preparar una lista de problemas GEMM en formato JSON describiendo transposición y disposición de A, B y C/D según sea necesario (por ejemplo, GEMMs FP16 con entrada/salida FP16 y cómputo FP32, HSH).
• Construir CUTLASS con la integración de heurísticas habilitada y especificar -DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE y -DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM para controlar cuántas configuraciones emite por problema.
• La etapa de compilación produce una lista CSV de pruebas que enumera todos los casos a probar para realizar el tuning sobre las configuraciones emitidas. Se puede usar código de benchmarking propio o cutlass_profiler para ejecutar las configuraciones automáticamente.
• Para resultados de profiling consistentes, ejecute con clocks bloqueados. La documentación ofrece guías sobre uso e integración. El ejemplo de ocho configuraciones para FP16/FP32 (HSH) ilustra el aspecto práctico: compilando y probando estos ocho kernels, se puede seleccionar el mejor kernel sin el costo de una búsqueda exhaustiva, manteniendo un rendimiento cercano al óptimo. NVIDIA Developer Blog

Puntos clave

• Las heurísticas entregan un conjunto pequeño y dirigido de configuraciones GEMM con alto potencial, en lugar de una búsqueda exhaustiva.
• La integración con CUTLASS 4.2 facilita un flujo de trabajo con listas JSON de problemas, CSVs de pruebas y la opción de cutlass_profiler.
• Los resultados muestran ganancias sustanciales de tiempo sin sacrificar rendimiento cercano al óptimo.
• Los tamaños de cluster estáticos pueden entregar beneficios adicionales frente a enfoques dinámicos en ciertos escenarios.
• nvMatmulHeuristics es una parte central de los heurísticos de cuBLAS y está disponible en acceso anticipado con integración CUTLASS. NVIDIA Developer Blog

FAQ

• Q: ¿Qué hace nvMatmulHeuristics? A: Analiza los parámetros del GEMM y las capacidades del hardware para predecir un conjunto pequeño de configuraciones de kernel de alto potencial.
• Q: ¿Cómo se usa con CUTLASS? A: Compile CUTLASS con heurísticas activadas y proporcione una lista de problemas en JSON, además del número de configuraciones por problema; la compilación genera un CSV utilizado por profiling o cutlass_profiler.
• Q: ¿Qué ganancias de rendimiento se pueden esperar? A: En los resultados reportados, 16 configuraciones pueden lograr aproximadamente el 96% del rendimiento pico en ~150 minutos, frente a >700 minutos para la búsqueda exhaustiva. Otros workloads muestran ganancias similares y reducción de tiempos de construcción. NVIDIA Developer Blog
• Q: ¿Está disponible públicamente nvMatmulHeuristics? A: Sí, se describe como parte de los heurísticos de cuBLAS y está disponible en acceso anticipado con integración en CUTLASS. NVIDIA Developer Blog

Referencias

• https://developer.nvidia.com/blog/improving-gemm-kernel-auto-tuning-efficiency-on-nvidia-gpus-with-heuristics-and-cutlass-4-2/
• La entrada de NVIDIA citada es la fuente principal para nvMatmulHeuristics, su integración con CUTLASS 4.2 y los resultados descritos. NVIDIA Developer Blog