FlashAttention 4: Atenção Mais Rápida e Eficiente em Memória para LLMs

Introdução

A atenção por produto escalar dimensionado (SDPA) domina o tempo de inferência e o consumo de energia dos grandes modelos de linguagem (LLMs): quase todas as execuções de operadores acontecem dentro de uma única primitiva. A atenção pega consultas (Q) e as multiplica pelas chaves (K). Em seguida, normaliza o resultado através de uma operação softmax antes de multiplicar pelos valores (V) para produzir a saída. É limitada pela memória: ler Q/K/V várias vezes da memória de alta largura de banda e depois escrever repetidamente blocos intermediários.

As famílias FlashAttention reduzem esse gargalo mantendo mais dados no chip por mais tempo, aproveitando os multiprocessadores de streaming (SMs) da GPU. Primeiramente disponibilizado como código aberto em 2022, o FlashAttention reestruturou a rotina de atenção para calcular o softmax em tempo real. Apenas algumas linhas de Q × K são armazenadas no chip a qualquer momento.

Este artigo serve como um mergulho aprofundado voltado para praticantes sobre o FlashAttention 4 (FA4). Vamos resumir o que mudou em relação às versões anteriores e fornecer algumas orientações sobre adoção e benchmarking. Nosso objetivo é fornecer aos engenheiros de infraestrutura de LLM, engenheiros de kernel e equipes de plataforma de ML informações suficientes para saber quando – e se – o FA4 deve ser usado em sua stack.

Principais Conclusões

• A4 foi projetado como um kernel de atenção Blackwell-first, destinado à implementação em GPUs SM100. Ele possui um pipeline de 5 estágios especializado por warp, permitindo maior sobreposição e reutilização no chip em comparação com as gerações anteriores do FlashAttention.
• A eficiência é aprimorada no softmax por meio de duas técnicas principais: executar o software exp2() diretamente nos núcleos CUDA para reduzir a contenção nos SFUs e o reescalonamento adaptativo online para evitar cálculos desnecessários de reescalonamento, preservando a estabilidade.
• Ainda existem alguns recursos ausentes: o FA4 é forward-first e atualmente não possui implementações para o backward pass, varlen e GQA/MQA, o que limita seu uso para treinamento e algumas arquiteturas de modelo.
• Adote o FA4 de forma pragmática: use-o primeiro para inferência Blackwell, coloque-o atrás de testes de benchmarking + verificação de correção + alternativas, e mantenha o FA3 (Hopper)/FA2 (Ampere/Ada) como os padrões estáveis para necessidades de produção mais amplas.

Por que os kernels de atenção ainda dominam o custo dos LLMs

Apesar dos altos FLOPS fornecidos pelas GPUs, os modelos transformer frequentemente são limitados pela largura de banda da memória, devido à atenção precisar carregar Q, K e V da memória global várias vezes. Para calcular a saída de atenção para cada token de consulta, cada um deve fazer o produto escalar com até milhares de tokens-chave, normalizar e reescalar esse resultado e, finalmente, buscar os valores. Implementações simples calculam toda a matriz Q×K e a armazenam na memória. Isso rapidamente ultrapassa o tamanho da memória interna, levando a um desempenho prejudicado por leituras externas à memória.

O FlashAttention mitiga esse problema ao dividir a matriz de atenção em blocos. Em vez de materializar Q×K, ele transmite pequenos blocos de Q e K para a memória compartilhada, calcula o softmax sobre o bloco e mantém apenas o máximo e a soma correntes para cada linha de consulta. Isso reduz significativamente o tráfego de memória, fazendo com que o kernel se torne limitado pelo cálculo.

FlashAttention Evolution em um Olhar

É útil entender como o FlashAttention evoluiu ao longo do tempo. Cada versão aproveitou novos recursos de hardware para melhorar a velocidade ou a eficiência. O FA4 é especificamente otimizado para GPUs da geração Blackwell. Aqui está uma tabela de cada versão do FlashAttention com as principais melhorias que trouxeram e seu impacto no desempenho:

Novidades no FlashAttention 4

FlashAttention 4 continua o paradigma de computação em blocos estabelecido pelas versões anteriores do FlashAttention: tensores Q/K/V são processados bloco a bloco para maximizar a localidade de dados na memória interna e minimizar o tráfego para a memória global. Instâncias individuais do kernel (chamadas de array cooperativo de threads, ou CTA) trabalham juntas para calcular um ou mais blocos de saída transmitindo os blocos correspondentes de Q, K e V da memória global, e então projetando-os para os resultados finais da atenção.

Os componentes internos do FA4 consistem em um cronograma altamente pipelineado que sobrepõe várias etapas do cálculo de atenção para melhorar a ocultação de latência e maximizar a utilização de recursos da GPU. Em contraste com o FlashAttention 3, que possui essencialmente uma sobreposição de carregamento/cálculo em 2 etapas, o FA4 divide o cálculo em vários conjuntos de etapas concorrentes, com diferentes partes do pipeline sendo gerenciadas por grupos de warps separados. Esses grupos de warps se especializam cada um em uma única responsabilidade durante o "ciclo de vida" de um bloco na função de atenção.

Este modelo de especialização de warp atribui warps de 32 threads a tarefas específicas, como:

• Movimentação de dados (loads): buscando blocos de consulta, chave e valor da memória global para o armazenamento mais rápido no chip.
• Estágios de computação: Realizando pontuações de atenção parciais usando hardware de multiplicação-acumulação de matrizes com núcleos tensor.
• Lógica de Softmax e normalização: Operações de exponenciais + normalização para os pesos de softmax da atenção.
• Lógica de reescalonamento e redução: aplicação de ajustes numéricos ou combinação de resultados parciais antes da acumulação.
• Epílogo/operações de armazenamento: escrevendo os blocos concluídos do tensor de saída de volta à memória global.

A intuição por trás do modelo de execução básico é manter pipelines assíncronas para cada grupo de warp, de forma que o escalonador de warps da GPU possa alternar rapidamente entre os grupos de warp à medida que os operandos de origem se tornam disponíveis.

Exponenciais de Software via Núcleos CUDA

Um truque novo e particularmente chamativo no FA4 é como ele calcula a função exponencial usada no softmax. Normalmente, os kernels aproveitam as SFUs dedicadas na GPU para realizar operações matemáticas caras, como exp(). No entanto, como há relativamente poucas SFUs por SM, elas podem facilmente se tornar um gargalo na fila.

O FA4 contorna essa limitação ao fazer com que núcleos CUDA normais realizem exponenciais aproximadas em software. Especificamente, o FA4 contém uma implementação personalizada de uma aproximação polinomial exp2(x) com precisão de hardware. O FA4 simula o cálculo da exponencial usando um polinômio cúbico, que pode então ser executado em paralelo em muitos núcleos CUDA, em vez de se concentrar em alguns SFUs.

Conclusão: nada de ficar esperando as SFUs calcularem o softmax. Ao lidar com exp() em núcleos de uso geral, o FA4 permite que o warp que calcula o softmax acompanhe melhor o ritmo do pipeline.

Reescalonamento Softmax Online Mais Inteligente (Novo no FA4)

Outra inovação no FA4 está relacionada a como a acumulação do softmax mantém a estabilidade numérica. Como todas as versões do FlashAttention fazem o softmax “online” – acumulando valores parciais de max() e somas à medida que eles percorrem a sequência (para evitar estouro). Versões anteriores atualizavam constantemente o fator de escala em execução (geralmente o maior logit até então). Elas faziam o reescalonamento constantemente – sempre que viam um novo valor máximo, elas reescalonavam para preservar a estabilidade do softmax. O FA4 introduz uma política de reescalonamento adaptativa. Ele só faz o reescalonamento se o novo máximo que você vê for significativamente maior que o anterior. Isso, segundo relatos, reduz a frequência das operações de reescalonamento em um fator de 10. Consideremos o seguinte diagrama:

O diagrama compara duas versões de reescalonamento usadas no FlashAttention, permitindo um cálculo de atenção estável e eficiente em termos de memória. A estratégia original reescala a saída de cada bloco se o valor máximo aparecer, garantindo que somas pequenas permaneçam precisas ao normalizar as contribuições passo a passo. Com o reescalonamento adaptativo, rastreamos máximos locais e globais, reescalando apenas quando ocorrem saltos ou quando os blocos crescem. Isso reduz a sobrecarga enquanto ainda previne estouros e permite uma soma numericamente estável de termos de pequeno valor.

Compatibilidade e status atual

Esta tabela fornece uma visão geral da compatibilidade esperada do FlashAttention 4 (FA4) e do que funciona hoje em dia em toda a stack:

A passagem direta do FA4 já está amplamente disponível (comprometida) na árvore de código-fonte pública. Existem posts abertos em projetos como o repositório FlashAttention do Lab, onde os usuários estão solicitando suporte completo para passagem reversa/varlen/group/GQA. Isso implica que continua sendo uma questão pública em aberto.

O suporte a sequências de comprimento variável (processamento eficiente de lotes de sequências com comprimentos diferentes sem preenchimento) e GQA (Grouped-Query Attention) ainda não é suportado na passagem reversa. GQA é um truque de atenção onde múltiplas cabeças de atenção usam as mesmas projeções de chave/valor (às vezes usado em LLMs para reduzir o uso de memória). Versões anteriores do FlashAttention suportavam configurações no estilo MQA/GQA. No entanto, o kernel principal de avanço do FA4 provavelmente assumirá que o mesmo número de cabeças é usado para cada consulta. Portanto, provavelmente ainda não suporta corretamente cabeças agrupadas. Modelos com atenção de consulta agrupada (como algumas variantes do Llama2) não serão suportados pelo FA4 até que isso seja implementado.

Integração ao ecossistema: O FA4 ainda não foi adotado por muitas frameworks/camadas de nível superior porque é novo. Um exemplo é o motor de serviço vLLM (altamente otimizado para inferência LLM), que atualmente não possui suporte ao FA4 – apenas o FlashAttention 2 e 3 são suportados. Selecionar FA4 no vLLM resultará em um erro de “versão não suportada” porque o seletor de versão deles atualmente só permite escolher entre as versões 2 ou 3.

Como Adotar o FlashAttention – Guia de Decisão

Use esta tabela para determinar qual versão do FlashAttention é a ideal para suas necessidades.

Perguntas Frequentes

Qual problema o FlashAttention 4 (FA4) está tentando resolver?

Reduz o pesado tráfego de memória HBM da atenção ao calcular a atenção bloco a bloco no chip, evitando a materialização da matriz completa QK⊤.

Qual é a principal mudança arquitetônica no FA4 em comparação com o FA3?

O FA4 utiliza um pipeline de ~5 estágios especializado em warp (carregar, calcular, softmax/normalizar, reescalar/reduzir, armazenar) para aumentar a sobreposição e o reaproveitamento no chip no Blackwell (SM10.x).

Por que o FA4 calcula exponenciais nos núcleos CUDA em vez de nos SFUs?

Como os SFUs são escassos e podem facilmente se tornar um gargalo ao calcular o softmax, o FA4 recorre a calcular uma aproximação por software de exp2() nos núcleos CUDA.

O que a “reescalonagem online adaptativa” altera no softmax?

Em vez de redimensionar sempre que um novo máximo aparece, o FA4 só redimensiona quando esse máximo aumenta de forma significativa. Isso reduz a sobrecarga de redimensionamento enquanto mantém a estabilidade numérica.

Quando um engenheiro deve adotar o FA4 na produção hoje?

Comece com Blackwell apenas com inferência, comprimento fixo, atenção padrão atrás de uma flag de recurso; mantenha cuDNN/SDPA (e FA3/FA2 em GPUs mais antigas) como alternativas até que o suporte a backward/varlen/GQA esteja maduro.

Conclusão

O FlashAttention 4 se baseia no objetivo do FlashAttention de reduzir as demandas de largura de banda de memória da atenção ao realizar mais trabalho de Q/K/V no chip. No entanto, ao contrário das versões anteriores, o FA4 faz isso com um pipeline de 5 estágios especializado em warp específico do Blackwell e duas melhorias importantes no softmax (usando exponenciais de software nos núcleos CUDA para aliviar gargalos do SFU; e usando reescalonamento adaptativo online para eliminar cálculos de estabilidade redundantes). O FA4 apresentou melhorias respeitáveis na taxa de processamento do forward-pass em comparação com o cuDNN em sequências longas, mas continua priorizando o forward e apenas para Blackwell. Além disso, o FA4 atualmente apresenta lacunas significativas de recursos (como ausência de backward, varlen, GQA, etc.), e o suporte do ecossistema ainda não está totalmente desenvolvido. A conclusão por enquanto é considerar os caminhos do FA4 para inferência Blackwell como uma possibilidade de alto potencial, que requer adoção após benchmarks rigorosos, verificações de correção e alternativas. Ao mesmo tempo, cargas de trabalho de treinamento e hardware mais amplo devem continuar a usar os caminhos de código maduros do FA3/FA2 para produção.

Referências

• FlashAttention-3: Atenção Rápida e Precisa com Assincronismo e Baixa Precisão
• FlashAttention 4: Quebrando a Barreira do Petaflop em Kernels de Atenção de GPU
• Como aplicar FA4 no B200?
• Suporte para passagem reversa FA4 no Blackwell

Obrigado por aprender com a comunidade Linux-Console.net.