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 orçamento energético dos modelos de linguagem de grande escala (LLMs): quase todas as execuções de operadores ocorrem dentro de um único primitivo. A atenção pega em queries (Q) e multiplica-as pelas keys (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: lê Q/K/V várias vezes da memória de alta largura de banda e depois escreve repetidamente os blocos intermédios.

As famílias FlashAttention reduzem este estrangulamento ao manter mais dados no chip por mais tempo, aproveitando os multiprocessadores de streaming (SMs) da GPU. Lançado 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 uma análise aprofundada orientada para profissionais sobre o FlashAttention 4 (FA4). Vamos resumir o que mudou em relação às versões anteriores e fornecer algumas orientações sobre adopção/benchmarking. O nosso objetivo é fornecer aos engenheiros de infraestruturas LLM, engenheiros de kernels e equipas de plataformas de ML informação suficiente para saber quando – e se – o FA4 deve ser usado na sua stack.

Principais Conclusões

• O A4 foi concebido como um núcleo de atenção Blackwell-first, destinado à implementação em GPUs SM100. Apresenta um pipeline de 5 etapas especializado para warp, permitindo um aumento da sobreposição e reutilização no chip em comparação com as gerações anteriores do FlashAttention.
• A eficiência é melhorada no softmax através de duas técnicas principais: executar o software exp2() diretamente nos núcleos CUDA para reduzir a contenção nos SFUs e o reescalamento adaptativo online para evitar cálculos de reescalamento desnecessários, mantendo a estabilidade.
• Ainda existem algumas funcionalidades em falta: o FA4 é forward-first e atualmente não tem implementações para o backward pass, varlen, e GQA/MQA, o que limita o seu uso para treino 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 corretude + alternativas, e mantenha o FA3 (Hopper)/FA2 (Ampere/Ada) como os padrões estáveis para necessidades de produção mais amplas.

Porque é que os kernels de atenção ainda dominam o custo dos LLM

Apesar dos elevados FLOPS fornecidos pelas GPUs, os modelos transformer são muitas vezes limitados pela largura de banda da memória, devido à atenção ter de carregar Q, K e V da memória global várias vezes. Para calcular a saída da atenção para cada token de consulta, cada um deve fazer produto interno com até milhares de tokens-chave, normalizar e reajustar este resultado e, finalmente, buscar os valores. Implementações simples calculam toda a matriz Q×K e armazenam-na na memória. Isto excede rapidamente o tamanho da memória no chip, levando a um desempenho prejudicado pelas leituras fora do chip.

O FlashAttention mitiga este 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 partilhada, calcula o softmax sobre o bloco e retém apenas o máximo e a soma em execução para cada linha de consulta. Isto reduz significativamente o tráfego de memória, fazendo com que o kernel se torne limitado pela computação.

FlashAttention Evolution de Relance

É útil compreender como o FlashAttention evoluiu ao longo do tempo. Cada versão tirou partido de novas funcionalidades 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 o seu impacto no desempenho:

Novidades no FlashAttention 4

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

Os componentes internos do FA4 consistem numa programação altamente pipelineada que sobrepõe várias fases do cálculo de atenção para melhorar a ocultação de latência e maximizar a utilização dos recursos da GPU. Em contraste com o FlashAttention 3, que possui essencialmente uma sobreposição de 2 fases de carregamento/cálculo, o FA4 divide o cálculo em vários conjuntos de fases concorrentes, com diferentes partes do pipeline sendo tratadas por grupos de warps separados. Estes grupos de warps especializam-se cada um numa única responsabilidade durante o “ciclo de vida” de um tile através da função de atenção.

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

• Movimento de dados (loads): buscar blocos de consulta, chave e valor da memória global para o armazenamento mais rápido no chip.
• Estágios de cálculo: Realização de pontuações de atenção parciais usando hardware de multiplicação-acumulação de matrizes com núcleos tensor.
• Softmax e lógica de normalização: Operações de exponenciais + normalização para os pesos softmax de 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: escrita dos blocos concluídos do tensor de saída de volta na memória global.

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

Exponenciais de Software através de Núcleos CUDA

Um truque novo 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 dispendiosas como exp(). No entanto, como existem relativamente poucas SFUs por SM, elas podem facilmente tornar-se um gargalo na fila.

O FA4 contorna esta 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 depois ser executado em paralelo através de muitos núcleos CUDA em vez de estar concentrado em alguns SFUs.

Conclusão: não é preciso esperar pelas SFUs para calcular 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 na FA4 está relacionada com a forma como a acumulação do softmax mantém a estabilidade numérica. Tal como todas as versões do FlashAttention fazem “softmax online” – acumulando valores máximos parciais e somas à medida que percorrem a sequência (para evitar overflow). As versões anteriores atualizavam constantemente o fator de escala em execução (geralmente o logit máximo até ao momento). Elas reajustavam constantemente – sempre que encontravam um novo valor máximo, reajustavam para preservar a estabilidade do softmax. A FA4 introduz uma política de reajuste adaptativo. Apenas reajusta se o novo máximo que encontrar for significativamente maior do que o anterior. Isto, segundo relatórios, reduz a frequência das operações de reajuste por um fator de 10. Consideremos o seguinte diagrama:

O diagrama compara duas versões de reescalamento 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 aparecer o valor máximo, garantindo que somas pequenas se mantenham precisas ao normalizar as contribuições passo a passo. Com o reescalamento adaptativo, monitorizamos os máximos locais e globais, reescalando apenas quando aparecem saltos ou quando os blocos aumentam de tamanho. Isto reduz a sobrecarga ao mesmo tempo que previne estouros e permite a soma numericamente estável de termos de valor pequeno.

Compatibilidade e estado 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 publicações abertas em projetos como o repositório FlashAttention do Lab, onde os utilizadores estão a pedir suporte completo para a passagem inversa/varlen/group/GQA. Isto implica que ainda é 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 backward. GQA é um truque de atenção onde múltiplas cabeças de atenção usam as mesmas projeções de chave/valor (por vezes usado em LLMs para reduzir o uso de memória). Versões anteriores do FlashAttention suportavam configurações ao estilo MQA/GQA. No entanto, o kernel principal de forward do FA4 provavelmente assumirá que o mesmo número de cabeças é usado para cada query. Portanto, provavelmente ainda não suporta corretamente cabeças agrupadas. Modelos com atenção agrupada por query (como algumas variantes do Llama2) não serão suportados pelo FA4 até que isto seja implementado.

Integração no 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 optimizado para inferência LLM), que atualmente não tem suporte ao FA4 – apenas o FlashAttention 2 e 3 são suportados. Selecionar FA4 no vLLM resultará num erro de “versão não suportada” porque o seu seletor de versões 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 mais adequada para as suas necessidades.

Perguntas Frequentes

Qual problema é que o FlashAttention 4 (FA4) está a tentar resolver?

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

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

O FA4 utiliza uma pipeline de cerca de 5 estágios especializada em warp (carregar, computar, softmax/normalizar, reescalar/reduzir, armazenar) para aumentar a sobreposição e a reutilização no chip no Blackwell (SM10.x).

Porque é que a FA4 calcula exponenciais nos núcleos CUDA em vez de nos SFUs?

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

O que altera o “redimensionamento online adaptativo” no softmax?

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

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

Comece com Blackwell apenas com inferência, comprimento fixo, atenção padrão atrás de uma flag de funcionalidade; 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 baseia-se no objetivo do FlashAttention de reduzir as exigências de largura de banda de memória da atenção, realizando 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 para Blackwell e duas melhorias importantes no softmax (utilizando exponenciais em software nos núcleos CUDA para aliviar os gargalos do SFU; e utilizando reescalamento adaptativo em tempo real para eliminar cálculos redundantes de estabilidade). O FA4 demonstrou melhorias respeitáveis na taxa de passagem para a frente em comparação com cuDNN em sequências longas, mas continua com preferência pelo processamento forward-first e é exclusivo para Blackwell. Além disso, o FA4 atualmente apresenta lacunas significativas de funcionalidades (como ausência de backward, varlen, GQA, etc.), e o suporte do ecossistema ainda não está totalmente desenvolvido. A conclusão por agora é considerar os caminhos do FA4 para inferência Blackwell como uma via potencialmente de alto retorno que requer adoção com base em benchmarking rigoroso, correção e alternativas. Ao mesmo tempo, as cargas de trabalho de treino e o hardware mais amplo devem continuar a usar os caminhos de código maduros FA3/FA2 para produção.

Referências

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

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