O que é o FlashAttention do ponto de vista da programação de kernels?

O FlashAttention é um kernel CUDA otimizado para calcular o mecanismo de atenção em transformadores, que acelera significativamente os cálculos e reduz o consumo de memória.

Do ponto de vista da programação de kernels, o FlashAttention implementa o cálculo de atenção usando uma gestão eficiente de memória e paralelismo na GPU. As principais características:

• Uso de blocos de memória partilhada para armazenar dados intermédios, o que reduz as chamadas à memória global lenta.
• Cálculo de atenção em threads com duplicação mínima de dados.
• Otimização do cálculo de softmax tendo em conta a estabilidade numérica.
• Processamento de sequências grandes sem necessidade de carregar todos os dados na memória ao mesmo tempo.

Isto permite realizar a atenção mais rapidamente e com menor consumo de memória em comparação com as implementações clássicas, o que é crítico para treinar modelos grandes.

Por exemplo, o kernel carrega blocos de chaves e consultas, calcula produtos escalares, aplica softmax e multiplica pelos valores, usando memória partilhada para acelerar e minimizar as chamadas à memória global.