FlashAttention-3 (PT)

FlashAttention-3 — é um algoritmo para otimizar o mecanismo de atenção (attention) em redes neurais transformer, desenvolvido para maximizar o uso dos recursos de hardware da arquitetura de GPU NVIDIA Hopper (H100). O algoritmo foi apresentado em 2024 por um grupo de pesquisadores das empresas Colfax Research, Meta, NVIDIA, Georgia Tech, Universidade de Princeton e Together AI. O trabalho foi aceito na conferência NeurIPS 2024 e destacado como spotlight.

O FlashAttention-3 é a terceira iteração na família de algoritmos, seguindo o FlashAttention (2022) e o FlashAttention-2 (2023). Seu principal objetivo é acelerar significativamente o treinamento e a inferência de grandes modelos de linguagem (LLMs), preservando ao mesmo tempo a precisão dos cálculos.

Introdução e histórico

O problema do mecanismo de atenção

O componente-chave dos transformers é o mecanismo de autoatenção (self-attention), no entanto, sua complexidade computacional e consumo de memória crescem de forma quadrática (O(n²)) com o aumento do comprimento da sequência de entrada (n). Isso cria um sério "gargalo", pois as GPUs modernas são otimizadas para multiplicações de matrizes rápidas, mas o cálculo de funções exponenciais (por exemplo, no Softmax) é ordens de magnitude mais lento. Além disso, em uma implementação ingênua, um grande tensor de atenção intermediário deve ser armazenado na memória da GPU, o que limita a escalabilidade dos modelos.

FlashAttention e FlashAttention-2

Para resolver esse problema, em 2022 foi proposto o FlashAttention, que reduziu o volume de acessos à lenta memória global (HBM) por meio de duas técnicas:

• Processamento em blocos (tiling): Os cálculos são divididos em blocos (tiles) que são processados na rápida memória on-chip (SRAM).
• Fusão de operações: Todas as operações (multiplicação de matriz, Softmax) são executadas em um único kernel da GPU sem gravar resultados intermediários na memória global.

Isso permitiu reduzir a complexidade de memória de quadrática para linear e acelerou os cálculos em 2 a 4 vezes.

Em 2023, foi apresentada uma versão aprimorada — FlashAttention-2, que otimizou a paralelização dos cálculos. Na arquitetura de GPU NVIDIA Ampere (A100), ela atingiu ~70% do desempenho teórico máximo. No entanto, na arquitetura mais recente NVIDIA Hopper (H100), sua eficiência foi significativamente menor — cerca de 35%. Isso ocorreu porque o algoritmo não utilizava os novos recursos de hardware da Hopper, o que motivou a criação do FlashAttention-3.

Novos recursos de hardware da GPU Hopper (H100)

A arquitetura NVIDIA Hopper introduziu uma série de novos recursos que o FlashAttention-3 utiliza para alcançar o máximo desempenho:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): Um novo tipo de instrução para os tensor cores, que realiza multiplicações de matriz com um aumento de desempenho de quase duas vezes em comparação com a arquitetura Ampere.
• TMA (Tensor Memory Accelerator): Um módulo de hardware que acelera a transferência de dados entre a memória global (HBM) e a memória compartilhada (shared memory). O TMA realiza cálculos de endereço automaticamente, descarregando os núcleos de computação.
• Formato FP8: Suporte de hardware para o formato de dados de ponto flutuante de 8 bits, que dobra o desempenho teórico em comparação com o FP16, mas traz o risco de perda de precisão devido à sua faixa dinâmica limitada.

Inovações técnicas do FlashAttention-3

O algoritmo implementa três métodos-chave de otimização, projetados especificamente para a arquitetura Hopper:

1. Execução assíncrona e especialização de warps

O FlashAttention-3 utiliza o princípio de warp specialization, no qual diferentes grupos de threads (warps) na GPU se especializam em tarefas distintas:

• Warps produtores (Producer warps): Carregam dados da memória global usando o TMA.
• Warps consumidores (Consumer warps): Realizam multiplicações de matriz nos tensor cores.

Graças à assincronia de hardware da Hopper, essas operações se sobrepõem no tempo. Enquanto um grupo de warps realiza os cálculos, outro carrega paralelamente os dados para o próximo bloco. Essa abordagem de pipeline (pipeline), organizada pelo princípio de "ping-pong" (ping-pong scheduling), permite ocultar as latências de operações lentas (como o Softmax) и utilizar ao máximo todos os módulos funcionais da GPU.

2. Minimização de operações de memória

O algoritmo mantém a ideologia de tiling das versões anteriores, mas utiliza ativamente o TMA para carregar de forma assíncrona os próximos blocos de dados em paralelo com os cálculos atuais. A transferência de dados da lenta HBM para a rápida SRAM é efetivamente realizada "na sombra" dos cálculos principais, fazendo com que a GPU passe menos tempo ociosa aguardando dados.

3. Baixa precisão (FP8) com redução do erro de quantização

A transição para FP8 dobra a velocidade, mas pode levar a uma perda significativa de precisão devido à quantização. Para combater isso, os desenvolvedores implementaram o método de processamento incoerente (incoherent processing). Sua essência é a seguinte:

1. Antes de calcular a atenção, os vetores de características (queries Q e keys K) são multiplicados por uma matriz ortogonal aleatória (por exemplo, uma matriz de Hadamard).
2. Essa transformação "espalha" os valores com magnitudes anomalamente grandes (outliers) por todas as coordenadas, uniformizando sua distribuição.
3. Depois disso, a quantização para FP8 é realizada, agora com um erro menor.
4. Como a transformação é ortogonal, ela não distorce o resultado final da atenção (QKᵀ), pois o efeito da matriz é anulado na multiplicação.

Essa técnica permitiu reduzir o erro no cálculo da atenção em FP8 em aproximadamente 2,6 vezes em comparação com a aplicação padrão de FP8 sem a transformação.

Desempenho e importância

A aplicação das técnicas mencionadas permitiu ao FlashAttention-3 alcançar uma superioridade significativa sobre as versões anteriores na GPU H100:

• Aceleração de 1,5 a 2 vezes em comparação com o FlashAttention-2.
• Alta utilização da GPU: Atinge ~75–85% do desempenho teórico máximo da H100.
• Taxa de transferência (throughput):
  • Até 740–840 TFLOPS para precisão meia (FP16/BF16).
  • Até 1,2–1,3 PFLOPS (petaflops) ao usar precisão de 8 bits (FP8).

A alta eficiência do FlashAttention-3 impacta diretamente o desenvolvimento e a aplicação de LLMs:

• Redução do tempo de treinamento: Uma aceleração de 75 a 100% na atenção reduz significativamente o tempo de treinamento dos modelos, que pode levar semanas ou meses.
• Aumento da janela de contexto: Os modelos podem processar eficientemente sequências mais longas (centenas de milhares de tokens), o que é importante para analisar grandes documentos ou códigos.
• Uso racional de recursos: Permite alcançar o mesmo desempenho com um número menor de GPUs ou obter maior velocidade no mesmo hardware, o que reduz o custo de implantação dos modelos.

Disponibilidade e integração

Os autores publicaram o código-fonte do FlashAttention-3 sob uma licença de código aberto no GitHub. Espera-se sua integração nos principais frameworks de aprendizado profundo, como PyTorch e as bibliotecas Hugging Face Transformers, o que tornará a tecnologia acessível a um amplo círculo de desenvolvedores e pesquisadores. As versões anteriores já se tornaram um padrão de fato na indústria, e o FlashAttention-3 provavelmente continuará essa tendência.

Ligações externas

• Repositório oficial do FlashAttention no GitHub
• Blog da Together AI com o anúncio do FlashAttention-3

Literatura

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.