Flash Attention

Flash Attention é um algoritmo otimizado para calcular o mecanismo de atenção no núcleo dos modelos transformadores. Introduzido por Tri Dao et al. da Stanford em um artigo de 2022, ele produz resultados numericamente idênticos à atenção padrão, mas reorganiza a computação para minimizar o movimento de dados entre a memória de banda larga (HBM) da GPU e a SRAM rápida no chip. É classificado como um algoritmo consciente de IO—seu design é explicitamente motivado pelo gargalo de largura de banda de memória do hardware moderno de GPU, em vez de reduzir a complexidade aritmética da operação de atenção em si.

A atenção padrão materializa a matriz de atenção N×N completa na HBM, onde N é o comprimento da sequência, resultando em uso de memória O(N²). Flash Attention usa uma estratégia de tiling: processa atenção em blocos pequenos o suficiente para caber em SRAM, aplicando o algoritmo softmax online numericamente estável para acumular resultados parciais incrementalmente. A matriz de atenção completa nunca é escrita na HBM. Flash Attention 2 (2023) melhorou o paralelismo na dimensão da sequência e reduziu operações de ponto flutuante não matmul. Flash Attention 3 (2024) introduziu especialização assíncrona de warp e suporte de precisão FP8 direcionado para GPUs NVIDIA H100 e classe Hopper, aproximando o throughput dos picos teóricos de hardware.

Atenção é o gargalo dominante de memória e computação no treinamento e inferência de transformadores em comprimentos de sequência longos. Flash Attention reduz a memória de O(N²) para O(N) e oferece acelerações de tempo de parede de 2–4x em relação a implementações não otimizadas, traduzindo-se diretamente em custos de treinamento mais baixos ou mais iterações experimentais por orçamento de computação. Esses ganhos se compõem em comprimentos de sequência mais longos, tornando Flash Attention uma infraestrutura essencial para modelos práticos de contexto longo.

A partir de 2026, Flash Attention é integrado nativamente no PyTorch via torch.nn.functional.scaled_dot_product_attention, Hugging Face Transformers, vLLM e praticamente todo grande framework de treinamento e inferência LLM. É uma dependência de treinamento para LLaMA 3, Mistral e a maioria dos outros modelos de ponta. Ring attention, que distribui tiling de Flash Attention em múltiplos dispositivos via comunicação all-to-all, estende seus benefícios a sequências que excedem a memória de GPU única, permitindo as janelas de contexto de milhões de tokens agora oferecidas comercialmente.

Exemplo

“Uma equipe de pesquisa treinando um transformador de 7 bilhões de parâmetros em sequências de 32.000 tokens muda da atenção padrão do PyTorch para Flash Attention 2, reduzindo o uso de memória GPU por passo em aproximadamente 60% e cortando o tempo de treinamento por passo em 3x, permitindo tamanhos de lote maiores e reduzindo pela metade o tempo total de parede de treinamento.