FlashAttention-3: Attention mais rápido e mais preciso com assincronia e baixa precisão
(together.ai)- O gargalo de Attention nos Transformers vem limitando a velocidade de treinamento e inferência de LLMs com contexto longo, e o FlashAttention-3 é uma nova versão que busca reduzir isso ainda mais usando recursos das GPUs Hopper
- O ponto central é usar a assincronia do Tensor Core e do TMA para sobrepor computação e movimentação de dados, além de intercalar matmul e softmax para reduzir o tempo ocioso da GPU
- No H100, o FlashAttention-2 aproveitava apenas 35% do FLOPS teórico máximo, mas o FlashAttention-3 chega a até 740 TFLOPS em FP16, alcançando 75% do valor teórico do H100
- Em baixa precisão FP8, outliers nas ativações de LLMs ampliam os erros, então ele reduz o erro em 2,6x em relação ao attention FP8 de referência com incoherent processing baseado em transformada de Hadamard
- O FlashAttention-3 foi publicado no GitHub, e quanto mais o custo de Attention cai, maior fica o potencial de melhorar o processamento de contextos mais longos e a eficiência de treinamento e inferência de LLMs
Objetivos e desempenho do FlashAttention-3
- Attention é a camada central do Transformer, mas se torna um gargalo principal em modelos de linguagem de grande porte e aplicações com contexto longo
- FlashAttention e FlashAttention-2 abriram caminho para acelerar Attention reduzindo leituras e escritas na memória da GPU, e hoje a maioria das bibliotecas usa essa abordagem para acelerar treinamento e inferência de Transformers
- Essa abordagem ajudou a ampliar o comprimento de contexto dos LLMs nos últimos 2 anos, de 2~4K no GPT-3 e OPT para 128K no GPT-4 e 1M no Llama 3
- O FlashAttention-2 ficava em 35% de aproveitamento do FLOPS teórico máximo na GPU H100, mas o FlashAttention-3 eleva isso usando novos recursos da GPU Hopper
- O desempenho divulgado do FlashAttention-3 é o seguinte
- 1,5~2,0x mais rápido que o FlashAttention-2 em FP16
- até 740 TFLOPS em FP16
- 75% de aproveitamento do FLOPS teórico máximo do H100
- quase 1,2 PFLOPS em FP8
- erro 2,6x menor que o attention FP8 de referência
Relembrando a abordagem do FlashAttention
- O FlashAttention acelera o cálculo de Attention reorganizando a ordem das operações e usando tiling e recomputação, além de reduzir o uso de memória em relação ao comprimento da sequência de quadrático para linear
- Ele carrega blocos de entrada da HBM para a SRAM, executa Attention nesses blocos e depois atualiza a saída na HBM
- Como não grava a grande matriz intermediária de Attention na HBM, reduz leituras e escritas de memória e pode atingir ganho real de 2~4x em tempo de execução
- Com tiling e reescalonamento de softmax em conjunto, é possível processar por blocos e ainda obter a saída correta sem aproximação
Recursos da GPU Hopper: WGMMA, TMA, FP8
- O FlashAttention-2 pode chegar a até 70% do FLOPS teórico máximo na GPU Ampere A100, mas não conseguia aproveitar bem os novos recursos da Hopper
- O FlashAttention-3 usa três recursos da Hopper
- WGMMA: recurso de multiplicação-acumulação matricial por warpgroup que usa os novos Tensor Cores da Hopper e tem throughput maior que o
mma.syncda Ampere - TMA: unidade dedicada de hardware que acelera a transferência de dados entre memória global e memória compartilhada, tratando cálculo de índices e predicação fora dos limites para reduzir o uso de registradores
- FP8: pode dobrar o throughput do Tensor Core em relação a FP16, mas como representa valores de ponto flutuante com menos bits, há um tradeoff com precisão
- WGMMA: recurso de multiplicação-acumulação matricial por warpgroup que usa os novos Tensor Cores da Hopper e tem throughput maior que o
- O FlashAttention-3 usa as abstrações do NVIDIA CUTLASS para aproveitar os recursos da Hopper
- Só de reescrever o FlashAttention para usar esses novos recursos, o desempenho do forward pass em FP16 sobe de cerca de 350 TFLOPS no FlashAttention-2 para 540~570 TFLOPS
Sobrepondo GEMM e softmax com assincronia
- As operações principais do Attention são os GEMMs entre Q-K e P-V, além do softmax
- Em aceleradores modernos, operações que não são matmul são muito mais lentas que matmul, e funções especiais como a exponencial do softmax são processadas em unidades separadas das de floating point multiply-add ou matrix multiply-add
- O H100 SXM5 entrega 989 TFLOPS em multiplicação matricial FP16, mas o throughput de funções especiais é de 3,9 TFLOPS, 256 vezes menor
- Com dimensão de head 128, mesmo que o FLOPS de matmul seja 512 vezes maior que o da exponencial, a exponencial ainda pode responder por 50% do tempo em relação ao matmul
- Em FP8, o FLOPS de matmul dobra, mas a velocidade da exponencial permanece igual, então executar matmul e softmax em paralelo se torna ainda mais importante
-
agendamento pingpong entre warpgroups
- O escalonador de warps da GPU já faz parte do overlap automaticamente ao executar outros warps enquanto alguns esperam pelo resultado do GEMM
- O FlashAttention-3 usa barreiras de sincronização para sobrepor de forma manual e mais eficiente o GEMM e o softmax de dois warpgroups
- o warpgroup 1 executa primeiro o GEMM1 de uma iteração e o GEMM0 da próxima iteração
- depois, enquanto o warpgroup 2 executa o GEMM, o warpgroup 1 processa o softmax
- esse agendamento pingpong esconde o softmax atrás do tempo de execução do GEMM do outro warpgroup
- na prática, o agendamento não fica tão limpo quanto no diagrama, mas eleva o throughput do forward pass de Attention em FP16 de cerca de 570 TFLOPS para 620 TFLOPS com head dimension 128 e sequence length 8K
-
overlap dentro do warpgroup
- Também é possível executar parte do softmax dentro de um único warpgroup enquanto o GEMM desse mesmo warpgroup está rodando
- Esse pipelining eleva o throughput do forward de Attention em FP16 de cerca de 620 TFLOPS para 640~660 TFLOPS
- Em troca, é preciso armazenar ao mesmo tempo o acumulador do GEMM e as entradas e saídas do softmax, o que aumenta a pressão sobre registradores
- No geral, essa técnica oferece um tradeoff vantajoso
FP8 de baixa precisão e incoherent processing
- Ativações de LLMs podem ter outliers com magnitude muito maior que a das demais features
- Outliers dificultam a quantização e ampliam bastante o erro de quantização
- O FlashAttention-3 usa incoherent processing, adotado em trabalhos de quantização como o QuIP
- Ele multiplica query e key por uma matriz ortogonal aleatória para dispersar os outliers e reduzir o erro de quantização
- Na implementação, é usada uma transformada de Hadamard com sinal aleatório
- Se a dimensão do head for
d, isso pode ser feito por head de attention em tempo O(d log d), e não O(d²) - Como a transformada de Hadamard é limitada por largura de banda de memória, ela pode ser fundida com operações anteriores também limitadas por largura de banda, como rotary embedding, sem custo adicional
- Se a dimensão do head for
- Em um experimento que gerou Q, K e V a partir de uma distribuição normal padrão e inseriu magnitudes grandes em 0,1% das entradas para simular outliers, o incoherent processing reduziu o erro de quantização em 2,6x
Benchmarks e estado de disponibilidade
- O FlashAttention-3 foi comparado não só com o FlashAttention-2, mas também com implementações em Triton e cuDNN que já usam os novos recursos de hardware da GPU Hopper
- Em FP16, ele mostra ganho de cerca de 1,6~1,8x sobre o FlashAttention-2
- Em FP8, chega a quase 1,2 PFLOPS
- O repositório GitHub do FlashAttention-3 foi publicado
- O artigo também pode ser consultado no mesmo repositório flash-attention
Otimizações restantes e integração futura
- O artigo inclui, além do que foi abordado no blog, otimizações como variable length sequence, persistent kernel e transpose in-kernel para FP8
- Projetar o algoritmo de acordo com o hardware de execução pode gerar grandes ganhos de eficiência e viabilizar novos recursos de modelo, como contexto mais longo
- Trabalhos futuros incluem otimizações para inferência de LLMs e a generalização das técnicas para outras arquiteturas de hardware
- Espera-se que o FlashAttention-3 seja integrado a futuras versões do PyTorch
1 comentários
Comentários do Hacker News
Pelos comentários no código, parece que Tri Dao já vinha trabalhando no FA3 desde abril de 2022, logo após o anúncio do Hopper/H100
É um pouco interessante que o código tenha levado mais de 2 anos para ser aberto hoje; talvez isso seja porque soluções melhores estejam a caminho
O histórico recente de artigos do Tri está mais inclinado para estruturas da linha SSM e Mamba. FlashAttention tem complexidade de tempo quadrática em relação ao comprimento da sequência, enquanto os algoritmos mais recentes são quase quadráticos, então não se limitam a fazer o mesmo cálculo de forma mais eficiente, mas reduzem bastante o volume total de computação
Dao e Gu mostraram este ano, em um artigo longo, que Mamba/SSM também pode ser formulado de forma favorável para aceleração com as mesmas operações primitivas de hardware das quais os Transformers se beneficiam
Se a SETH for provada ou refutada, o problema P vs NP também será resolvido, então é difícil esperar que isso aconteça tão cedo
O ponto principal é se um caso de uso específico consegue arcar com esse custo
Fico curioso sobre o quanto o algoritmo FlashAttention é amarrado ao hardware
Por exemplo, neste anúncio dizem que ele aproveita os recursos assíncronos da GPU H100; então isso parece significar que placas fora da linha H não conseguem esse ganho de velocidade
Além disso, a biblioteca FlashAttention de uso real exige CUDA, mas o algoritmo aparentemente foi portado para Metal[^0]. Se o algoritmo for algo mais próximo de uma função pura, não deveria ser possível implementá-lo em qualquer GPU/framework de machine learning?
[0]: https://github.com/philipturner/metal-flash-attention
As versões mais recentes variam conforme o nível de abstração. ThunderKittens[0] oferece, de forma parecida com o que o texto menciona, ganhos de velocidade de cerca de 1,3x a 2x sobre o FA2, mantendo aplicação relativamente geral em GPUs
Cada novo hardware pode ter recursos específicos que extraem desempenho adicional. Em geral, os vendors adotam recursos que os coloquem à frente, mas, assim como já acontece no CUDA, as APIs e bibliotecas acabam ficando fragmentadas
[0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
Queria perguntar ao pessoal de compiladores: será que um compilador conseguiria descobrir sozinho uma otimização como o FlashAttention? TVM e tinygrad parecem estar indo nessa direção, mas é difícil acreditar que isso seja possível
Só que o custo é alto, então seria preciso armazenar em cache os resultados dessa busca
A otimização com e-graphs parece se encaixar bem nessa área. Mas isso exigiria uma grande mudança de paradigma na forma como os passes de otimização são tratados, então quase não foi implantado fora de algumas ferramentas de nicho. Por exemplo, isso não combina bem com o grafo de chamadas tradicional; para aplicar e-graphs atravessando blocos básicos e indo para fora/entre loops
for, seria necessário mudar bastante o fluxo de controle, ebreakereturntambém não são suportadosNão sei até onde chegou o estado da arte em otimização de compiladores em termos de layout de dados e maximização da utilização do processador
Vi uma vez um vídeo sobre otimização dizendo que pequenas otimizações até aumentavam a velocidade, mas o impacto era pequeno comparado à variação de desempenho causada pelas diferenças de layout de memória criadas por essa otimização ou até mesmo por mudanças aleatórias
A apresentação focava mais em distinguir o sinal do ruído, mas esse próprio ruído já é um indício de que compiladores não lidam muito bem nem com formas bem mais simples do problema discutido aqui
Só a arquitetura de CPU e memória, em que cache e padrões de acesso afetam a velocidade, já é complexa; se somarmos a arquitetura de GPU, isso parece ser uma área bastante inexplorada
Talvez um dia isso seja possível. Como estamos falando de IA, também surge a pergunta se uma IA suficientemente inteligente conseguiria fazer isso, mas depende do que significa “suficientemente”
Dá para imaginar isso como um teste de altíssimo nível para modelos de IA: dar algo como micrograd e pedir que façam algo mais rápido que o torch mantendo a mesma interface. Ainda estamos longe disso, mas seria interessante se se tornasse possível
TVM faz sentido. Estritamente falando ele faz outra coisa, mas fica bem perto disso
Já o tinygrad, não entendi por que deu essa impressão
Se alguém quiser portar isso para ROCm / AMD MI300x, entre em contato em hello@hotaisle.xyz. Nunca enviarei spam
Posso doar tempo de computação para esse trabalho
Na 4090, talvez seja possível obter ganho de velocidade usando a implementação Triton de atenção FP8: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
Esta é uma das melhorias mais importantes de toda a IA. Ela permite fazer mais coisas, mais rápido, com o mesmo hardware, e traz um ganho para quase todos os usuários de IA com praticamente nenhum trade-off
Queria entender por que o FlashAttention fica cerca de 5 vezes mais lento ao usar mascaramento variável do que sem isso. Sem um bom suporte a masking, o efeito da otimização praticamente desaparece
Seria ótimo se algum especialista pudesse responder a algumas perguntas :)
O FlashAttention é um substituto direto, tipo drop-in, para a operação de attention em LLMs? Dá para usar em qualquer lugar onde a operação de “attention” é usada, ou é preciso treinar o LLM separadamente para usar FA?
Qual é a relação do FA com estratégias como GQA e sliding window attention? São conceitos ortogonais entre si, ou cada estratégia precisa de uma implementação separada de FA?
Recentemente o llama.cpp adicionou suporte a FlashAttention; isso significa que ele passou a usar algo como os kernels CUDA fornecidos pelo FlashAttention?
Por fim, neste texto é feita uma comparação entre FlashAttention e Triton. Triton não seria algo como uma camada de abstração? Não daria para implementar FA em Triton? A expressão “FlashAttention vs Triton” não ficou muito clara para mim
https://github.com/Dao-AILab/flash-attention
FlashAttention é uma forma de calcular a parte
Softmax(QK^T)Vda attention, enquanto GQA é uma forma de calcular as matrizes Q, K e V. Já a sliding window attention, embora eu tenha menos certeza, funciona alterando a máscara de attention que controla quais chaves cada consulta pode enxergarNão usei o llama.cpp, mas a explicação de que ele começou a aproveitar kernels CUDA parece, em linhas gerais, correta
A última pergunta está se referindo a uma implementação anterior de FlashAttention escrita em Triton
Fiquei curioso porque o texto diz que operações como sigmoid são muito lentas
LLMs modernos usam bastante funções de ativação com sigmoid ou Softmax, como SiLU, Swish e SOLU
O ReLU tem menos perda de desempenho? Se sim, talvez não valha mais a pena voltar ao velho e bom ReLU?