FlashAttention-2: Attention mais rápida com melhor paralelização e divisão de trabalho

• Surgiram modelos de linguagem com contextos mais longos, como GPT-4(32k), MPT(65k) e Calude(100k)
• Expandir o comprimento de contexto dos transformadores é difícil porque as exigências de tempo de execução e memória aumentam à quarta potência
• O FlashAttention, lançado no ano passado, passou a ser usado em vários lugares por reduzir o uso de memória e aumentar a velocidade da attention
• No lançamento, ele já era de 2 a 4 vezes mais rápido, mas ainda havia espaço para melhorias. Em comparação com operações de multiplicação de matrizes otimizadas (GEMM), ainda não era tão rápido e atingia apenas 25-40% do FLOPs/s máximo teórico (máximo de 124 TFLOPs/s na GPU A100)
• O FlashAttention-2 é 2 vezes mais rápido que a versão anterior e oferece desempenho de até 230 TFLOP/s na GPU A100
• No treinamento de modelos de linguagem do tipo GPT, chegou a até 225 TFLOPS (72% de utilização de FLOPs do modelo)
• O algoritmo foi ajustado para reduzir FLOPs que não são de multiplicação de matrizes
• Melhor paralelização e mudança na forma de dividir o trabalho em cada bloco de threads
• Expansão da quantidade de Head Dimensions de 128 para 256