¿Qué es FlashAttention desde el punto de vista de la programación de kernels?

FlashAttention es un kernel CUDA optimizado para calcular el mecanismo de atención en transformadores, que acelera significativamente los cálculos y reduce el consumo de memoria.

Desde el punto de vista de la programación de kernels, FlashAttention implementa el cálculo de atención utilizando una gestión eficiente de la memoria y el paralelismo en la GPU. Las principales características:

• Uso de bloques de memoria compartida para almacenar datos intermedios, lo que reduce las llamadas a la memoria global lenta.
• Cálculo de atención en hilos con duplicación mínima de datos.
• Optimización del cálculo de softmax teniendo en cuenta la estabilidad numérica.
• Procesamiento de secuencias grandes sin necesidad de cargar todos los datos en memoria al mismo tiempo.

Esto permite realizar la atención más rápido y con menor consumo de memoria en comparación con las implementaciones clásicas, lo cual es crítico para entrenar modelos grandes.

Por ejemplo, el kernel carga bloques de claves y consultas, calcula productos escalares, aplica softmax y multiplica por los valores, usando memoria compartida para acelerar y minimizar las llamadas a la memoria global.