🤖 machine learning

Hardware-Aware FP4 FlashAttention-4

Este artigo introduz o Hardware-Aware FP4 FlashAttention-4, um método que supera as limitações dos tensor cores FP4 da Blackwell ao empregar Direct-P para inferência não causal e um caminho causal com gradientes FP8, alcançando até 2,13× mais velocidade de throughput de forward e 1,14× mais velocidade de atualizações de treinamento em GPU única, enquanto evita os problemas de divergência observados no treinamento MXFP4.

Autores originais: Robert Hu

Publicado 2026-09-04✓ Author reviewed ⓘ

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Robert Hu

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

No mundo da inteligência artificial, os modelos mais poderosos dependem de um mecanismo chamado "atenção" para compreender o contexto. Imagine um leitor examinando um documento longo; a atenção permite que o modelo foque nas palavras mais relevantes em uma frase enquanto ignora o restante, conectando ideias através de vastas distâncias. Para fazer isso, o computador realiza uma série massiva de cálculos, comparando cada palavra com todas as outras para determinar suas relações. Durante anos, esses cálculos foram realizados com números de alta precisão, semelhante ao uso de uma régua com marcações minúsculas e precisas para medir um edifício. Isso garante que o modelo aprenda corretamente, mas é lento e consome enormes quantidades de energia. À medida que os modelos crescem, a necessidade de velocidade torna-se crítica. Pesquisadores desenvolveram recentemente chips capazes de usar números muito menores e mais grosseiros — valores de ponto flutuante de quatro bits — para realizar essas comparações muito mais rápido. No entanto, simplesmente mudar para esses números menores não é suficiente; o software que gerencia o fluxo de dados também deve mudar, ou os ganhos de velocidade desaparecem.

Um pesquisador da Graphcore abordou esse gargalo específico com um novo método que eles chamam de Direct-P. O trabalho deles foca na passagem "forward" (direta) da atenção, onde o modelo processa informações para gerar uma resposta, e na passagem "backward" (reversa), onde ele aprende com seus erros. O pesquisador descobriu que, embora os novos chips possam multiplicar números incrivelmente rápido, a etapa intermediária — converter pontuações brutas em probabilidades — estava atrasando tudo. Era como ter uma linha de montagem super-rápida que parava constantemente porque um trabalhador tinha que medir cuidadosamente cada peça antes de passá-la adiante. O pesquisador descobriu que a forma padrão de lidar com essa conversão, que envolvia escalonamento e arredondamento complexos, criava um congestionamento que negava os benefícios de velocidade do novo hardware.

Para resolver isso, o pesquisador redesenhou o processo de conversão para ser direto e simplificado. Em vez de calcular uma probabilidade precisa e depois arredondá-la, seu método mapeia as pontuações brutas diretamente para os códigos específicos que o hardware utiliza. Isso elimina as etapas intermediárias que causavam atrasos. Quando testaram essa abordagem nos chips de última geração para data centers, os resultados foram impressionantes. Para certas formas comuns de dados, o novo método processou informações mais do que duas vezes mais rápido do que o padrão anterior, que dependia de números de maior precisão. Ele alcançou isso mantendo a saída precisa o suficiente para tarefas complexas como geração de vídeo e compreensão de linguagem. Em testes envolvendo um modelo de geração de vídeo, o novo método foi quase duas vezes mais rápido que a abordagem padrão, produzindo resultados que, embora finitos e utilizáveis, mostraram desvio mensurável e pontuações de similaridade mais baixas em comparação com a versão mais lenta e precisa.

No entanto, a história não é apenas sobre velocidade; é também sobre o que acontece quando o modelo tenta aprender. O pesquisador explorou se poderiam usar esses números de baixa precisão ultra-rápidos para todo o processo de treinamento, incluindo a passagem "backward", onde o modelo atualiza seu conhecimento. Eles descobriram que, embora a passagem "forward" pudesse rodar em velocidade total, a passagem "backward" exigia um compromisso cuidadoso. Quando tentaram usar o formato de quatro bits mais rápido para os valores de probabilidade durante o treinamento, o processo de aprendizado tornou-se instável e o modelo falhou em melhorar. Os números tornaram-se muito grosseiros, fazendo com que o sinal de aprendizado quebrasse. Consequentemente, o pesquisador determinou que, para o treinamento permanecer estável, eles deveriam usar um formato de oito bits ligeiramente mais preciso para os valores de probabilidade, mesmo que o restante do cálculo utilize o formato de quatro bits mais rápido. Essa abordagem híbrida permitiu acelerar todo o ciclo de treinamento em cerca de 14 por cento em um único chip poderoso, um ganho significativo para modelos de grande escala.

O pesquisador também observou os limites físicos dos próprios chips de computador. Ele descobriu que a velocidade do cálculo não era mais o problema principal; a questão era como os dados eram armazenados e movimentados dentro do chip. O chip possui uma área de memória pequena e ultra-rápida onde mantém os números enquanto trabalha neles. O pesquisador descobriu que esse espaço de memória estava completamente cheio, não deixando espaço para sobrepor diferentes estágios do cálculo. Era como uma cozinha onde o balcão está tão lotado de ingredientes que o chef não pode começar a picar o próximo vegetal até que o atual seja finalizado, mesmo que a faca seja incrivelmente afiada. Como o espaço de memória estava totalmente ocupado, o chip não conseguia trabalhar em múltiplas etapas ao mesmo tempo, o que limitava o quanto o processo inteiro poderia se tornar mais rápido.

Este trabalho destaca uma mudança crucial em como pensamos em tornar a inteligência artificial mais rápida. Não basta simplesmente construir calculadoras mais rápidas; todo o fluxo de trabalho deve ser redesenhado para corresponder às capacidades do hardware. O pesquisador mostrou que, ao mudar a forma como as probabilidades são calculadas e ao gerenciar cuidadosamente o fluxo de dados, eles poderiam desbloquear melhorias massivas de velocidade. No entanto, eles também provaram que existem limites rígidos. As restrições de memória do chip significam que, mesmo com a aritmética mais rápida possível, há um teto para o quanto de sobreposição é possível. O caminho a seguir, sugerem eles, não reside apenas em aritmética mais rápida, mas em formas mais inteligentes de organizar os dados para que os recursos do chip nunca fiquem esperando. Esse equilíbrio entre velocidade bruta e gerenciamento cuidadoso de dados é o que permitirá que a próxima geração de inteligência artificial opere de forma eficiente.

=== RESUMO TÉCNICO ===

Resumo Técnico: Hardware-Aware FP4 FlashAttention-4

Declaração do Problema

As GPUs NVIDIA Blackwell introduzem núcleos tensoriais de ponto flutuante de 4 bits (FP4) que aceleram significativamente a multiplicação de matrizes. No entanto, nos mecanismos de atenção, o ganho de velocidade é frequentemente limitado não pelos produtos de matrizes ( $QK^T$ e $PV$ ), mas pela etapa intermediária de softmax. Esta etapa envolve redução de pontuação, avaliação exponencial, construção de probabilidade e escalonamento, que devem ser concluídos antes que o segundo produto de matriz possa começar.

As tentativas existentes de utilizar FP4 para todo o bloco de atenção enfrentam dois obstáculos primários:

Tempo (Timing): A construção do tensor de probabilidade ( $P$ ) a partir das pontuações ( $S$ ) cria uma dependência de caminho crítico que atrasa o produto $PV$ .
Faixa Numérica (Numerical Range): Os formatos FP4 (especificamente NVFP4 e MXFP4) utilizam um payload de 4 bits (E2M1) com faixa dinâmica limitada. As probabilidades do softmax padrão abrangem uma faixa que é inadequada para este payload, levando a underflow, "escalas zero" (onde blocos inteiros desaparecem) ou perda significativa de precisão se não forem gerenciadas cuidadosamente.

Além disso, no treinamento causal, a passagem de retropropagação (backward pass) normalmente requer a reconstrução de probabilidades a partir de estados salvos. Abordagens que salvam pontuações de precisão total ou as recomputam em BF16 anulam o potencial de aceleração do treinamento de baixa precisão.

Metodologia: Direct-P e Causal Backward

O artigo propõe o Direct-P, um método para permitir inferência direta em FP4 total e um caminho causal de retropropagação quantizado em hardware Blackwell (GB200/B300).

1. Direct-P para Inferência de Forward

O Direct-P modifica o intervalo entre a conclusão do fragmento de pontuação $QK^T$ e a emissão da instrução $PV$ . Ele retém o pipeline existente de dois queries da HAO AI Lab e o protocolo de propriedade de Tensor Memory (TMEM), mas otimiza a geração de probabilidade:

Mapeamento Direto (Direct Mapping): Em vez de computar uma exponencial de alta precisão e depois arredondar, o Direct-P trata a construção da probabilidade como um problema de classificação de código. Ele mapeia pontuações log-normalizadas diretamente para as 8 magnitudes E2M1 não negativas usando um classificador afim (ou exponenciais nativas seletivas). Isso encurta o caminho crítico.
Normalização Consistente: O denominador (normalizador) para o softmax é acumulado diretamente a partir dos códigos de probabilidade FP4 arredondados e das escalas de bloco que serão consumidos pela instrução $PV$ . Isso garante que o numerador e o denominador descrevam o mesmo operador aproximado, evitando discrepâncias entre somas de alta precisão e payloads de baixa precisão.
Proteção de Faixa (Range Guarding): Para camadas com logits extremos (ex: >500), um mecanismo de âncora amostrada evita underflows subnormais sem exigir um re-escaneamento completo das pontuações.
Seleção de Formato: O método utiliza NVFP4 para $Q$ e $K$ (devido ao escalonamento fino) e MXFP4 para $P$ e $V$ (devido às suas escalas de bloco de potência de dois que se alinham bem com os fragmentos de produtor $N_{32}$ ).

2. Causal Backward Quantizado

Para o treinamento, o artigo introduz um caminho onde a passagem de forward salva o payload quantizado de $Q/K$ , as escalas e o normalizador log-sum-exp (LSE).

Reconstrução: O backward pass reconstrói as probabilidades diretamente a partir desses operandos quantizados salvos, em vez de criar um caminho de pontuação BF16 separado.
Formatos de Gradiente: Os epílogos de projeção publicam visões orientadas por linha/coluna em FP8 de $Q, K, V$ para alimentar os produtos de matriz de gradiente.
Restrição de Estabilidade: Embora o caminho de forward utilize MXFP4 para $P/V$ , experimentos de treinamento distribuído revelaram que as trajetórias de MXFP4 P/V divergem. Consequentemente, a rota de treinamento retida utiliza FP8 para $P$ e $V$ no backward pass para garantir estabilidade, mantendo $Q/K$ em NVFP4.

Contribuições Principais

Implementação Direct-P: Um método que mapeia pontuações diretamente para códigos de probabilidade MXFP4 e normaliza usando os mesmos valores representados, alcançando até 2,13× de throughput de forward em BF16 no GB200 para formas favoráveis, com um ganho de velocidade de média geométrica de 2,02× através da grade D128 testada.
Caminho de Treinamento Causal Quantizado: Um backward pass que reconstrói probabilidades diretamente dos estados de quantização de forward salvos, acelerando uma atualização completa de um modelo de 8 bilhões de parâmetros em um único GPU em até 1,14×.
Perfilamento de Hardware: Identificação da propriedade de Tensor Memory (TMEM) como o limite primário para o overlap. O estudo demonstra que ter apenas núcleos tensoriais mais rápidos é insuficiente se o layout de armazenamento (512 colunas em D128) impedir que a próxima operação $QK$ se sobreponha à fase $PV$ atual.

Resultados

Desempenho de Inferência de Forward

Throughput: No GB200, a rota "fast" do Direct-P atinge um ganho de velocidade de média geométrica de 2,02× sobre a baseline de BF16 da HAO, atingindo o pico de 2998 TFLOP/s (S32768/H24).
Precisão: O Direct-P troca parte da precisão por velocidade. Em testes de softmax Gaussiano, a rota "fast" possui uma similaridade de cosseno média de ~0,944 contra BF16, comparado a ~0,990 para a rota NV/FP8 da HAO.
Impacto no Modelo:
- ViT/BERT: Em avaliações de entrada fixa, a rota "fast" mantém a precisão da tarefa (ex: a precisão top-1 do ViT S4096 iguala-se ao BF16 em 88,5%) apesar da menor similaridade de cosseno no nível do operador.
- Wan Video Diffusion: O método permanece finito e funcional para 20 passos de difusão. No entanto, o erro acumula-se ao longo dos passos; o modelo de 14B em 20 passos mostra um L2-relativo de ~0,53 em comparação ao BF16, indicando deriva (drift), mas geração funcional.

Desempenho de Treinamento

Single-GPU: A rota quantizada (projeções NVFP4 + P/V em FP8) acelera a subcamada de atenção completa em 1,25× e a atualização completa do modelo de 8B em 1,14× com batch size 4.
Treinamento Distribuído: Em um experimento correspondente de 64 GPUs sobre 100 bilhões de tokens, a rota P/V em FP8 alcançou uma melhoria de throughput de 1,112× sobre o BF16.
Estabilidade: O artigo observa explicitamente que as trajetórias de MXFP4 P/V divergiram em configurações distribuídas. A rota de treinamento bem-sucedida depende do uso de FP8 para o produto $P/V$ , o que significa que o caminho de treinamento não é "puro" FP4 de ponta a ponta.

Significância e Alegações

O artigo afirma que o Direct-P torna o forward da atenção em FP4 total praticável ao encurtar o caminho crítico entre os produtos de matriz. Demonstra que, embora os núcleos de matriz FP4 sejam mais rápidos, o gargalo mudou para movimentação de dados, construção de probabilidade e propriedade de memória.

Principais Conclusões:

Limites de Hardware: A barreira primária para maiores ganhos de velocidade não é o throughput aritmético, mas a propriedade de TMEM. O layout atual D128 consome todas as 512 colunas de TMEM, impedindo que a próxima operação $QK$ se sobreponha ao final da fase $PV$ .
Paradas de Prontidão (Readiness Stalls): Baixa atividade de tensor-core é frequentemente um sintoma de prontidão de operandos (esperando pela construção da probabilidade) em vez de falta de capacidade de computação.
Trade-offs de Treinamento: Embora o forward de inferência em FP4 seja viável, o treinamento causal estável atualmente exige uma concessão ao FP8 para os operandos de probabilidade/valor para evitar divergência. O artigo sugere que formatos futuros como Unsigned E5M3 (UE5M3) podem permitir o treinamento totalmente em FP4, mas isso não é avaliado aqui.

O trabalho conclui que aritmética de matriz mais rápida sozinha é insuficiente; alcançar maior eficiência requer mudanças arquiteturais que permitam janelas de overlap alocáveis maiores (ex: destinos de pontuação adicionais) juntamente com algoritmos de baixa precisão otimizados.