DeepGEMM: kernels GEMM FP8 limpios y eficientes mediante escalado fino

 (github.com/deepseek-ai)
2 puntos por GN⁺ 2025-02-27 | 1 comentarios | Compartir por WhatsApp
  • DeepGEMM es una biblioteca de kernels tensor core de alto rendimiento que reúne en una sola base de código CUDA primitivas de cálculo clave de los LLM modernos, como GEMM, fused MoE, MQA scoring y HyperConnection.
  • Todos los kernels se compilan en runtime mediante un módulo JIT liviano; no requiere compilación CUDA durante la instalación y exige C++20, CUDA Toolkit, PyTorch y CUTLASS 4.0 o superior.
  • Usa algunos conceptos de CUTLASS y CuTe, pero no depende de forma pesada de templates ni de estructuras algebraicas; está diseñada para facilitar el aprendizaje de la optimización de kernels NVIDIA GPU con un número limitado de funciones de kernel centrales.
  • Su cobertura incluye FP8, FP4, BF16 GEMM, grouped GEMM, kernels de MQA logits para DeepSeek v3.2 y Mega MoE que solapa comunicación y cómputo, con distintas restricciones de layout de memoria en SM90 y SM100.
  • A pesar de su diseño liviano, apunta a igualar o superar el rendimiento de bibliotecas ajustadas por expertos en distintas shapes de matrices, e incluye una actualización que alcanzó hasta 1550 TFLOPS en H800.

Propósito y diseño de DeepGEMM

  • DeepGEMM es una biblioteca de kernels tensor core que integra en una sola base de código CUDA las principales primitivas de cálculo usadas en los modelos de lenguaje grandes modernos.
    • GEMM: FP8, FP4, BF16
    • fused MoE con comunicación solapada: Mega MoE
    • MQA scoring para lightning indexer
    • HyperConnection(HC)
  • Todos los kernels se compilan en runtime mediante un módulo Just-In-Time(JIT) liviano.
    • No requiere compilación CUDA durante el proceso de instalación.
  • Usa algunos conceptos de CUTLASS y CuTe.
    • Pero no depende en gran medida de templates pesados ni de estructuras algebraicas.
    • Mantiene la base de código simple al limitar la cantidad de funciones de kernel centrales.
  • Afirma que, pese a su diseño liviano, ofrece rendimiento equivalente o superior al de bibliotecas ajustadas por expertos en varias shapes de matrices.

Actualizaciones principales

  • La actualización del 16 de abril de 2026 incluye Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, compilación JIT más rápida y más.
    • Detalles en #304
    • Benchmark de Mega MoE en #316
  • La actualización del 28 de septiembre de 2025 agrega un kernel de scoring weighted ReLU MQA logits para el lightning indexer de DeepSeek v3.2.
  • La actualización del 20 de julio de 2025 da soporte tanto a SM90 como a SM100 y refactoriza todo con un módulo JIT CPP de bajo overhead de CPU.
    • NVRTC y la optimización SASS posterior a la compilación están desactivados.
    • NVRTC queda marcado como soporte futuro.
    • Como NVCC 12.9 realiza automáticamente FFMA interleaving, la optimización posterior a la compilación ya no se soporta.
    • Detalles en #112
  • La actualización del 14 de mayo de 2025 agrega kernels de weight gradient para dense y MoE backward.
    • Detalles en #95
  • La actualización del 7 de mayo de 2025 ofrece hasta 10 veces más velocidad de compilación con soporte para NVRTC.
    • Se puede activar con DG_JIT_USE_NVRTC=1.
    • En algunos casos puede haber pérdida de rendimiento.
    • Detalles en #94
  • La actualización del 18 de abril de 2025 alcanza hasta 1550 TFLOPS en H800.

Requisitos y flujo de instalación

  • El entorno de ejecución requiere una GPU con arquitectura NVIDIA SM90 o SM100.
  • Los requisitos de software son los siguientes:
    • Python 3.8 o superior
    • Compilador con soporte para C++20
    • CUDA Toolkit
      • SM90: CUDA 12.3 o superior
      • Se recomienda fuertemente CUDA 12.9 o superior para obtener el máximo rendimiento
      • SM100: CUDA 12.9 o superior
    • PyTorch 2.1 o superior
    • CUTLASS 4.0 o superior
    • Biblioteca {fmt}
  • En el entorno de desarrollo, se clona el repositorio incluyendo los submodules y luego se ejecuta develop.sh para enlazar los includes necesarios y compilar el módulo JIT CPP.
  • La instalación consiste en ejecutar install.sh y luego importar deep_gemm desde un proyecto Python.

Interfaz GEMM y restricciones de layout

  • La convención de nombres de los kernels GEMM de DeepGEMM es D = C + A @ B.
  • El layout de shape de entrada se basa en NT.
    • fp8_gemm_nt ejecuta D = C + A @ B.T.
  • La implementación para SM90 solo soporta layout de memoria NT.
    • Corresponde a la combinación row-major, col-major.
  • La implementación para SM100 soporta todos los layouts de memoria NT, TN, NN, TT.
  • En ambas arquitecturas, el scaling factor del LHS debe estar alineado para TMA y en layout transpuesto.
    • SM90 requiere el scaling factor en formato FP32.
    • SM100 requiere el formato packed UE8M0, empaquetando 4 UE8M0 en un torch.int.
  • Operaciones como transponer la entrada o convertir a FP8 deben ser manejadas por separado por el usuario.
    • La biblioteca ofrece funciones utilitarias simples de PyTorch, pero pueden ser lentas.
    • El foco principal es la optimización de kernels GEMM.

Dense y Grouped GEMM

  • El FP8 GEMM básico non-grouped usa las funciones fp8_gemm_{nt, nn, tn, tt}.
  • El grouped GEMM con contiguous layout, a diferencia del grouped GEMM tradicional de CUTLASS, solo agrupa el eje M.
    • N y K deben ser fijos.
    • Es un diseño pensado para situaciones en modelos MoE donde los experts comparten la misma shape.
  • En el forward pass de entrenamiento o en el prefilling de inferencia, la cantidad de tokens procesados por cada expert puede variar.
    • La forma de concatenar esos tokens en un único tensor se denomina contiguous layout.
    • Cada segmento de expert debe estar alineado al GEMM M block size.
    • El criterio de alineación se consulta con get_mk_alignment_for_contiguous_layout().
  • También se ofrece una API de grouped por eje K para MoE weight backward.
    • M y N deben ser fijos.
    • La función relacionada es k_grouped_fp8_gemm_tn_contiguous.
  • En la etapa de inference decoding, cuando CUDA graph está habilitado y la CPU no puede conocer la cantidad de tokens por expert, se soporta masked grouped GEMM.
    • Si se proporciona un mask tensor, el kernel calcula solo las regiones válidas.
    • La función es m_grouped_fp8_gemm_nt_masked.
    • Hay un ejemplo que usa como entrada la salida del kernel de baja latencia de DeepEP.

Kernel MQA para DeepSeek v3.2 Indexer

  • La familia de kernels MQA V3.2 ofrece una versión non-paged y una versión paged.
    • La non-paged es para prefilling.
    • La paged es para decoding.
  • fp8_mqa_logits recibe 6 entradas.
    • q: tensor E4M3, shape [seq_len, num_heads, head_dim]
    • kv: tensor E4M3 y float scaling factor
      • La shape del tensor es [seq_len_kv, head_dim]
      • La shape del scaling factor es [seq_len_kv]
    • weights: tensor float, shape [seq_len, num_heads]
    • cu_seq_len_k_start, cu_seq_len_k_end: tensor int, shape [seq_len]
    • clean_logits: si se deben limpiar los logits no llenados con -inf
  • La shape del tensor de salida es [seq_len, seq_len_kv] y representa logits token-to-token.
  • Cada token q i recorre los tokens kv j desde cu_seq_len_k_start[i] hasta antes de cu_seq_len_k_end[i].
    • Multiplica kv_j por el scaling factor.
    • Calcula valores por head con q[i, :, :] @ kv_j.
    • Después de aplicar ReLU, multiplica por weights[i, :] y suma para crear un scalar logit.
  • La función de la versión paged es fp8_paged_mqa_logits.

Mega MoE

  • Mega MoE fusiona varias etapas de MoE en un solo mega-kernel.
    • EP dispatch
    • linear 1, FP8xFP4
    • SwiGLU
    • linear 2, FP8xFP4
    • EP combine
  • Mega MoE solapa la comunicación NVLink con el cómputo tensor core.
  • Para ejecutarlo se necesita un multi-process launch que use symmetric memory.
  • El flujo de uso es el siguiente:
    • Asignar un symmetric memory buffer con deep_gemm.get_symm_buffer_for_mega_moe.
      • Requiere PyTorch 2.9 o superior.
    • Convertir los weights, incluidos FP4 y UE8M0 SF, al layout requerido con deep_gemm.transform_weights_for_mega_moe.
    • Antes de llamar, copiar input, scaling factor, top-k index y top-k weight al buffer.
    • Ejecutar el kernel fused mega MoE con deep_gemm.fp8_fp4_mega_moe.
  • El setup multi-process completo y el ejemplo de benchmarking están en tests/test_mega_moe.py.

Utilidades y variables de entorno

  • Las principales funciones utilitarias controlan recursos de ejecución, alineación, compilación JIT y conversión de scaling factors.
    • deep_gemm.set_num_sms / get_num_sms: establecer y consultar el número máximo de SM a usar
    • deep_gemm.set_tc_util / get_tc_util: establecer y consultar la tasa aproximada de utilización de tensor cores
    • deep_gemm.set_pdl / get_pdl: activar y desactivar Programmatic Dependent Launch(PDL)
    • deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout: establecer y consultar la alineación M/K a nivel de grupo del contiguous layout
    • deep_gemm.transform_sf_into_required_layout: convertir scaling factors al layout requerido
    • deep_gemm.get_tma_aligned_size: consultar el tamaño de alineación TMA necesario
  • Las variables de entorno relacionadas con JIT controlan la salida de debug, la ubicación de caché, la selección de compilador y las opciones de profiling.
    • DG_JIT_DEBUG: imprimir información de debug de JIT
    • DG_PRINT_CONFIGS: imprimir la config seleccionada por shape
    • DG_JIT_CACHE_DIR: directorio de caché de kernels compilados; el valor predeterminado es $HOME/.deep_gemm
    • DG_JIT_USE_NVRTC: usar NVRTC en lugar de NVCC; permite compilación rápida, pero en algunos casos puede tener menor rendimiento
    • DG_JIT_NVCC_COMPILER: ruta del compilador NVCC
    • DG_JIT_CPP_STANDARD: versión del estándar C++; el valor predeterminado es 20
  • También se ofrecen variables de entorno para debug y profiling.
    • DG_JIT_DUMP_ASM, DG_JIT_DUMP_PTX, DG_JIT_DUMP_SASS: volcar salida PTX y SASS
    • DG_JIT_WITH_LINEINFO: incluir información de líneas de source para herramientas de profiling
    • DG_COMM_KERNEL_DEBUG: inicializar el symmetric buffer en 0 antes de llamar a Mega MoE
    • DG_USE_NVIDIA_TOOLS: omitir el profiling interno al ejecutar herramientas externas de NVIDIA
  • Las opciones de build controlan la instalación y la forma de cargar kernels.
    • DG_SKIP_CUDA_BUILD: omitir el build de la extensión CUDA durante la instalación
    • DG_FORCE_BUILD: forzar build local en lugar de descargar un wheel pre-built
    • DG_JIT_USE_RUNTIME_API: usar CUDA Runtime API para cargar kernels; requiere CUDA runtime 12.8 o superior

Licencia y cita

  • El repositorio DeepGEMM está publicado bajo la MIT License.
  • El proyecto afirma estar inspirado en CUTLASS.
  • El título de la cita es DeepGEMM: clean and efficient BLAS kernel library on GPU

Lecturas relacionadas

1 comentarios

 
GN⁺ 2025-02-27
Opiniones de Hacker News

  • El intercalado de FFMA SASS se ve realmente asombroso
    Al ver que el rendimiento de los kernels FP8 de CUTLASS mejoró entre NVCC 12.2 y 12.3 y comparar el SASS compilado, parece que encontraron que un bit se había invertido con un patrón de intercalado en varias instrucciones FADD; tomando como referencia una implementación open source de un ensamblador CUDA, descubrieron que ese bit era el bit yield, que hace que el warp actual ceda el turno para que se ejecute otro warp
    Lo impresionante es que, aprovechando eso, crearon un script para modificar las instrucciones FFMA del binario compilado y, como cuando un warp cede no se puede reutilizar el registro, también invirtieron el bit reuse junto con el bit yield, logrando superponer mejor las instrucciones MMA con las instrucciones FFMA de promoción en GEMM FP8 con escalado fino, y en algunos casos aumentar el rendimiento más de 10%

    • Por lo que leí en otros lados, en la optimización de operaciones matriciales donde el rendimiento es crítico, este tipo de enfoque suele ser bastante típico
      Pero en este problema específico parece que otras empresas de IA todavía no habían sentido la necesidad de aplicarlo, y al final es probable que todos lleguen a puntos similares
    • Scott Gray ya había descubierto exactamente esto y más en Maxwell en 2015, y desde entonces muchas otras personas lo han tratado bastante

  • Casos como este muestran muy bien lo lejos que están los compiladores actuales de extraer el rendimiento del hardware usando solo código de alto nivel
    Me pregunto qué haría falta para que las técnicas tradicionales de compiladores o agentes de optimización basados en IA produzcan resultados así

    • Parece que haría falta una enorme cantidad de prueba y error dentro de un bucle de retroalimentación con aprendizaje por refuerzo

  • Las cifras de mejora de velocidad reportadas son contra su propia línea base basada en CUTLASS
    Me pregunto si alguien hizo una comparación directa de rendimiento con cuBLAS
    Los resultados de CUTLASS GEMM que he visto hasta ahora estaban más o menos dentro de 10% frente a cuBLAS; si se mantiene la mejora de 2x a 2.5x que menciona el paper, sería realmente impresionante

    • Normalmente evito FP8 y prefiero I8, pero esta pregunta me dio curiosidad por saber qué tan bien sale cuBLAS
      Para empezar, cuBLAS necesita la API extendida cuBLASLt para manejar operaciones de precisión mixta como FP8
      Además, no soporta combinaciones de tipos que parecerían adecuadas, como E5M2 x E5M2 en A x B, aunque sí soporta E5M2 x E4M3; y siguen las restricciones, por ejemplo que en Ampere, Hopper y Blackwell la matriz A siempre debe estar en layout transpuesto
      Integré un benchmark de FP8 cuBLASLt en mi repositorio "Less Slow C++" <https://github.com/ashvardanian/less_slow.cpp> y lo agregué a la lista de benchmarks existentes de cuBLAS y CUDA/PTX escritos a mano
      Lo estoy ejecutando en una GPU H200, que debería tener el mismo rendimiento que una H100, y en entradas cuadradas el throughput alcanza un pico de alrededor de 1.35 Peta-ops
      256 dio unos 2.68T/s, 512 20.49T/s, 1024 144.23T/s, 2048 665.68T/s, 4096 1.26P/s, 8192 1.34P/s y 16384 1.23P/s; eso es alrededor de 67% de la cifra que NVIDIA promociona para GEMM denso <https://resources.nvidia.com/en-us-data-center-overview-mc/e...>
    • Escuché que con CUTLASS se puede obtener mejor rendimiento que con cuBLAS
      Pensé que la línea base elegía lo mejor entre cuBLAS y CUTLASS

  • Este tipo de open source muestra muy bien el objetivo de la industria de lograr eficiencia
    Sin embargo, los beneficios de este software probablemente lleguen más a las grandes empresas que sirven modelos a gran escala, es decir, posibles competidores de DeepSeek, que a la comunidad open source general que quiere aprender, experimentar o servir modelos en hardware de consumo

    • Si la eficiencia mejora, al final puede traducirse en hardware más barato para todos, incluso para DeepSeek

  • No estoy muy seguro de que optimizar hacia precisiones cada vez más bajas sea bueno a largo plazo
    Esto implica que los modelos en realidad son bastante dispersos, y aunque ahora pueda ser así, creo que es más probable que se deba a que hay malas ideas mezcladas en la forma de entrenamiento, no a que necesariamente tengan que ser tan dispersos por naturaleza

    • Mientras funcione la dispersión gratis, hay que disfrutarla
      Hacer que solo se puedan entrenar modelos muy buenos con mayor precisión es un problema de investigación, mientras que el entrenamiento y la inferencia de baja precisión son problemas de ingeniería
      Venimos haciendo este tipo de cosas desde la época de las CNN, al menos desde hace 9 años, y creo que todavía quedan algunos años más
    • Como las funciones de activación descartan bastante rango dinámico de los números de punto flotante, parece bastante claro que probablemente no sea útil reservar un rango amplio para zonas de activación que ya están saturadas

  • Esto podría volverse irrelevante por MXFP, el soporte nativo de microescalado de Blackwell
    En Hopper, en cierto modo, implementaron eso manualmente con una granularidad más gruesa, pero usando coeficientes de escalado FP32

    • Exacto
      Una demostración pública de alta calidad como esta muestra muy bien dónde está el foso de $NVDA
      Las GPU de propósito general son muy flexibles, así que permiten programar muchas tareas que el proveedor de hardware no pensó desde el principio, pero que tienen bastante sentido
      Sin embargo, si uno predice que el futuro convergerá cada vez más hacia soporte dedicado en hardware y que desaparecerá el margen para estas optimizaciones de software, entonces el llamado foso de CUDA se derrumba
      Para seguir en este juego, NVIDIA está derribando su propio foso :p

  • Wow, tiene licencia MIT
    Ojalá las grandes empresas adopten este modelo de colaboración open source

  • Sigo preguntándome por qué existen instrucciones no documentadas
    Aunque no sean perfectamente estables, me parece que lo correcto sería ponerlas a disposición de los usuarios
    Este tipo de cosas seguramente está documentado internamente, pero no entiendo por qué no lo hacen público
    La seguridad basada en la oscuridad no funciona, y la competencia de todos modos hace ingeniería inversa de todo

    • Probablemente sea parecido a la razón por la que en lo que nosotros hacemos también terminan apareciendo partes no documentadas
      Puede ser por falta de tiempo, o porque no quieren insinuar soporte para funciones inestables o experimentales
      Si el daño solo llega, más o menos, al equipo de al lado, también es mucho más fácil cambiarlo
    • Puede que la premisa de que “este tipo de cosas está documentado internamente” ni siquiera sea cierta
      Probablemente solo aparezca en algún documento de diseño de arquitectura o especificación, y es natural que no quieran compartir ese tipo de documentos

  • Honestamente, es algo que está más allá de mi alcance de uso y de mi comprensión
    Aun así, me parece realmente valioso y refrescante que compartan estos hallazgos y mejoras para que todos puedan beneficiarse

    • FFMA significa Fused Floating-point Multiply-Add, y es una instrucción básica de GPU que realiza D = A*B + C en una sola operación
      Es muy importante en la multiplicación de matrices y en las cargas de trabajo de deep learning
      En SASS de NVIDIA, la instrucción FFMA se codifica como una instrucción de 64 o 128 bits y tiene varios bits de control que determinan su comportamiento exacto
      Cuando el bit yield está activado, le indica al planificador de warps que, después de esta instrucción, el warp actual puede ceder la ejecución, y el hardware puede ejecutar otro warp para ocultar la latencia
      La GPU obtiene alto rendimiento mediante paralelismo masivo, y si un warp se detiene, por ejemplo esperando memoria, otros warps pueden avanzar
      El bit reuse indica si un registro fuente puede reutilizarse en la operación inmediatamente siguiente, y cuando el bit yield está activado debe estar necesariamente desactivado
      Esto se debe a que, si el warp cede, puede que no sea el siguiente warp en ejecutarse y otros warps pueden cambiar el estado del archivo de registros, por lo que el hardware no puede garantizar que los valores de los registros se mantengan más allá del yield
      Si se activan los bits yield en un patrón alternado dentro de las instrucciones FFMA, el compilador crea puntos explícitos de planificación donde otros warps pueden avanzar, y para mantener la corrección también debe borrar los bits reuse de esas instrucciones
      Este cambio ayuda especialmente a solapar las instrucciones MMA, que son el núcleo de la multiplicación de matrices, con las instrucciones FFMA de promoción, que convierten FP8 para acumularlo con mayor precisión
      FP8 GEMM normalmente requiere convertir a una precisión más alta para la acumulación y luego volver a convertir, lo que genera FFMA adicionales; esto reduce los requisitos de ancho de banda de memoria, pero crea un patrón de cálculo complejo que mezcla operaciones de promoción y degradación
      “Escalado fino” parece referirse a gestionar cuidadosamente la precisión en varios puntos del cálculo
      La manipulación del bit yield permite intercalar mejor las operaciones de cálculo y de conversión de formato, haciendo un uso más eficiente de las unidades de ejecución de la GPU; sin esta optimización, el planificador de warps podría no encontrar oportunidades naturales de cambio y los recursos de cómputo podrían aprovecharse menos