FlashAttention-2: Attention más rápida gracias a una mejor paralelización y división del trabajo

• Han aparecido modelos de lenguaje con contextos más largos, como GPT-4(32k), MPT(65k) y Claude(100k)
• Extender la longitud de contexto de los transformers es difícil porque los requisitos de runtime y memoria aumentan a la cuarta potencia
• FlashAttention, lanzado el año pasado, se usa en muchos lugares porque reduce el uso de memoria y aumenta la velocidad de attention
• En el momento de su lanzamiento ya era entre 2 y 4 veces más rápido, pero todavía había margen de mejora. Seguía siendo más lento que una multiplicación de matrices optimizada (GEMM) y solo alcanzaba entre el 25% y el 40% de los FLOPs/s máximos teóricos (máximo de 124 TFLOPs/s en una GPU A100)
• FlashAttention-2 es 2 veces más rápido que la versión anterior y ofrece hasta 230 TFLOP/s de rendimiento en una GPU A100
• En el entrenamiento de modelos de lenguaje tipo GPT, llegó hasta 225 TFLOPS (72% de utilización de FLOPs del modelo)
• Se ajustó el algoritmo para reducir los FLOPs no relacionados con matmul
• Mejor paralelización y cambios en la forma de dividir el trabajo dentro de cada bloque de hilos
• Se amplió la cantidad de Head Dimensions de 128 a 256