FlashMLA: optimización de los núcleos de descodificación MLA para las GPU Hopper (Día 1 de la Semana del Código Abierto de DeepSeek)

Introducción general

FlashMLA es un programa desarrollado por la DeepSeek AI ha desarrollado un eficaz kernel de descodificación MLA (Multi-head Latent Attention) optimizado para las GPU de arquitectura Hopper de NVIDIA, diseñado para mejorar el rendimiento del procesamiento de secuencias de longitud variable. El proyecto se ha publicado en código abierto en GitHub y está a disposición de los desarrolladores de forma gratuita. Admite cálculo de precisión BF16 y caché KV paginada (tamaño de bloque 64), y ofrece un buen rendimiento en el H800 SXM5, con hasta 3000 GB/s de ancho de banda en configuraciones de uso intensivo de memoria y hasta 580 TFLOPS en configuraciones de uso intensivo de cálculo. FlashMLA se inspira en FlashAttention 2&3 y el proyecto Cutlass, y es apto para entornos de producción desde el primer momento. DeepSeek AI ha demostrado su innovación en tecnología de IA a través de este proyecto de código abierto, que ha atraído mucha atención.

Lista de funciones

• Decodificación MLA eficienteOptimizado para las GPU Hopper para acelerar significativamente el procesamiento de secuencias de longitud variable.
• Admite la precisión BF16Utiliza operaciones de coma flotante de media precisión para mejorar la eficiencia computacional manteniendo la precisión.
• Caché KV de paginaciónEl mecanismo de paginación con un tamaño de bloque de 64 se utiliza para gestionar eficazmente la memoria y mejorar el rendimiento de la inferencia.
• Alto rendimientoGPU H800 : Proporciona hasta 3000 GB/s de ancho de banda de memoria y 580 TFLOPS de potencia de cálculo en la GPU H800.
• código abiertoSe proporciona el código fuente completo para que los desarrolladores puedan realizar modificaciones e integraciones personalizadas.

Utilizar la ayuda

Proceso de instalación

FlashMLA es un proyecto de código abierto basado en GitHub , antes de utilizarlo es necesario asegurarse de que el entorno cumple los requisitos y completar la instalación . Aquí están los pasos detallados:

1. Preparación medioambiental

• sistema operativoSoporte para sistemas Linux (se recomienda Ubuntu 20.04 o superior).
• requisitos de hardware: Requiere una GPU con arquitectura Hopper de NVIDIA (como la H800 SXM5).
• dependencia del software::
  • CUDA 12.6 o superior (consulta las instrucciones de instalación en el sitio web de NVIDIA).
  • PyTorch 2.0 o superior (recomendado mediante pip install torch (Instalación).
  • Python 3.8 o superior.
• Herramientas de inspecciónAsegúrate de que Git está instalado para descargar código de GitHub.

2. Descargar el código fuente

1. Abra un terminal e introduzca el siguiente comando para clonar el repositorio FlashMLA:

   git clone https://github.com/deepseek-ai/FlashMLA.git

2. Vaya al catálogo de proyectos:

   cd FlashMLA
   

3. Instalación de dependencias

El proyecto depende de PyTorch y CUDA, que pueden instalarse con los siguientes comandos:

pip install -r requirements.txt

En caso contrario requirements.txt es sencillo asegurarse de que PyTorch está instalado:

pip install torch torchvision

Compruebe que CUDA está disponible:

python -c "import torch; print(torch.cuda.is_available())"

exportaciones True Indica que la configuración del entorno se ha realizado correctamente.

4. Compilación y pruebas

FlashMLA proporciona plug-ins CUDA precompilados, pero asegúrese de que coinciden con su versión CUDA local:

1. Vaya al directorio fuente y ejecute el script de compilación (si existe):

   python setup.py install
   

2. Compruebe que la instalación es correcta y ejecute el código de ejemplo:

   python example.py
   

Si no se informa de ningún error, la instalación se ha completado.

Cómo utilizarlo

La característica principal de FlashMLA es proporcionar un soporte eficiente de decodificación MLA para tareas de inferencia de modelos de IA. Estos son los pasos:

Función 1: Cargar y ejecutar FlashMLA

1. Módulo de importación::
   Introducción de las funciones básicas de FlashMLA en scripts de Python:

   from flash_mla import get_mla_metadata, flash_mla_with_kvcache
   

2. Preparación de la introducción de datos::
   • cache_seqlensDefine la longitud de secuencia de la caché KV.
   • q_i: Tensor de consulta.
   • kvcache_iDatos en caché KV.
   • block_tableTabla de bloques para la caché de paginación.
3. Obtener metadatos::

   tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv)
   

4. decodificación en marcha::

   o_i, lse_i = flash_mla_with_kvcache(q_i, kvcache_i, block_table, cache_seqlens, dv, tile_scheduler_metadata, num_splits, causal=True)
   

   exportaciones o_i es el resultado de la descodificación.lse_i es el valor de la suma logarítmica.

Función 2: Optimizar el tratamiento de secuencias de longitud variable

• tomaFlashMLA: cuando se trata de secuencias de entrada de longitud dinámica, FlashMLA reduce la huella de memoria paginando la caché KV.
• plataforma::
  1. Configurar los parámetros de paginación: el tamaño de bloque está fijado en 64 y puede ajustarse mediante el ajuste de los parámetros cache_seqlens Controla la longitud de la secuencia.
  2. Especificar en tiempo de ejecución causal=True(c) Garantizar que el mecanismo de atención causal está en funcionamiento.
• efecto3000 GB/s de ancho de banda de memoria en el H800 para tareas de inferencia a gran escala.

Función 3: Pruebas de rendimiento

• Métodos de ensayo::
  1. Edite el script de ejemplo (por ejemplo example.py), aumentando el tamaño de los datos de entrada.
  2. Utilice el siguiente código para registrar el rendimiento:

     import time
     start = time.time()
     # 运行解码代码
     o_i, lse_i = flash_mla_with_kvcache(...)
     print(f"耗时: {time.time() - start} 秒")
     

• Resultados esperadosCasi 3000 GB/s para tareas intensivas de memoria y 580 TFLOPS para tareas intensivas de cálculo.

advertencia

• compatibilidad de hardware: Sólo se admiten GPUs Hopper, se recomienda H800 o equivalente.
• Consejos de depuración: Si encuentras errores de CUDA, comprueba si la versión coincide o busca ayuda de la comunidad en GitHub Issues.
• entorno de producciónIntegración directa en los procesos de inferencia de modelos existentes, garantizando que los formatos de los datos de entrada sean coherentes con los requisitos de FlashMLA.

Con los pasos anteriores, los usuarios pueden empezar a utilizar FlashMLA rápidamente y disfrutar de la mejora de rendimiento que aporta su eficaz descodificación. El código completo y la documentación se encuentran en el repositorio de GitHub, y se recomienda ajustar los parámetros en función de los requisitos reales del proyecto.