FlashAttention-4: como Together AI acelerou atenção na GPU Blackwell

FlashAttention-4 — é um redesenho do algoritmo de atenção especificamente para GPUs modernas, onde os núcleos tensores crescem em desempenho significativamente mais rápido do que a memória e outros recursos.

Por que a abordagem padrão não funciona mais

À primeira vista, o desempenho de atenção (attention) é controlado pela velocidade de duas multiplicações matriciais: S = Q × K^T e O = P × V. Mas a análise da Blackwell B200 mostra algo inesperado: o gargalo não está nos núcleos tensores, mas em blocos funcionais especiais (SFU) para calcular a exponencial no softmax (forward pass) e no tráfego de memória compartilhada (backward pass).

De Hopper (H100) para Blackwell (B200), o desempenho dos núcleos tensores BF16 aumentou de 1 para 2.25 petaflops, enquanto a quantidade de SFU e a largura de banda da memória compartilhada permaneceram as mesmas. Essa assimetria quebra a otimização padrão — não se pode simplesmente assumir que os núcleos determinam completamente o desempenho. Na realidade, eles se atrapalham mutuamente.

Como Together AI resolveu o problema

O grupo de pesquisa junto com engenheiros da NVIDIA, Meta e Princeton propôs três ideias-chave:

• Novo pipelining — pipelines de software que maximalmente sobrepõem o trabalho dos núcleos tensores, SFU e memória sem ociosidade.
• Imitação da exponencial através de polinômios — em vez do lento bloco SFU, o forward pass calcula a exponencial em unidades FMA rápidas (fused multiply-add).
• TMEM e 2-CTA MMA — uso da nova tensor memory (256 KB por SM) mais um modo onde dois blocos de threads trabalham em uma operação matricial, reduzindo o tráfego de memória compartilhada.

Capacidades do Blackwell que tornaram isso possível

Blackwell adicionou especificamente alguns recursos para tais otimizações:

Tensor memory (TMEM) — armazenamento rápido no-chip (256 KB por cada SM), interconectado diretamente aos núcleos tensores. Resultados intermediários podem permanecer em TMEM sem ir para a mais lenta memória compartilhada, o que reduz drasticamente a latência de acesso.

Núcleos tensores assíncronos de 5ª geração — cada núcleo é iniciado por um único thread e acumula resultados em TMEM. O tile máximo para BF16 é 128×256×16 (aproximadamente 2 vezes maior que em Hopper), permitindo pipelining mais profundo sem overflow de registradores.

2-CTA MMA — novo modo onde dois blocos de threads trabalham simultaneamente em uma operação matricial. Isso reduz a carga na memória compartilhada pela metade e diminui o número de operações atômicas.

Números e resultados

FlashAttention-4 na Blackwell B200 com BF16 atinge 1605 TFLOPs/s (71% de utilização). Isso é 1.3× mais rápido que cuDNN 9.13 e 2.7× mais rápido que Triton. Para contexto: 1605 petaflops — é quase metade do desempenho de pico do Blackwell, mas extraído do complexo núcleo de atenção.

O que isso significa

FlashAttention-4 mostra como trabalhar na era do escalonamento assimétrico de GPUs — não refinar parafusos em algoritmos antigos, mas redesenhá-los conjuntamente com as novas capacidades do hardware. O resultado parece um feito, mas é apenas o começo da adaptação à nova realidade do hardware.

*Meta é reconhecida como organização extremista e proibida na RF.