FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Guia de leitura

Problema e motivação

Os Transformers possuem complexidade quadrática de tempo e memória em relação ao comprimento da sequência $N$. Embora métodos de atenção aproximada reduzam a contagem teórica de FLOPs, muitas vezes eles falham em proporcionar acelerações no tempo de execução real porque as computações em GPUs modernas são limitadas principalmente por transferências de memória (IO), em vez de operações aritméticas.

Na atenção padrão, o cálculo de $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$ e $\mathbf{O}=\mathbf{PV}$ grava matrizes intermediárias de tamanho $N\times N$ na lenta High Bandwidth Memory (HBM) e as lê repetidamente. Isso incorre em $\Theta (Nd+N^2)$ acessos à HBM e uma pegada de memória de $O(N^2)$, limitando o tamanho prático das janelas de contexto.

Ideias centrais

• Consciência de IO (IO-Awareness): Levar em consideração explicitamente a hierarquia de memória da GPU: a SRAM on-chip, que é rápida e pequena (por exemplo, 192 KB/SM com largura de banda de ~19 TB/s na A100), versus a HBM, grande e lenta (40–80 GB com largura de banda de 1,5–2,0 TB/s).
• Tiling (Cálculo do Softmax em Blocos): Dividir a entrada $\mathbf{Q},\mathbf{K},\mathbf{V}$ em blocos que caibam inteiramente na SRAM. Ao rastrear as estatísticas cumulativas de máximo por linha e normalizadores do softmax, a redução do softmax é calculada incrementalmente entre os blocos sem materializar a matriz $N\times N$ completa na HBM.
• Recomputação no Backward Pass: Evitar o armazenamento da matriz de atenção intermediária $N\times N$ para a retropropagação. Em vez disso, armazena-se apenas a saída $\mathbf{O}$ e as estatísticas do softmax $(\ell ,m)$, recomputando $\mathbf{S}$ e $\mathbf{P}$ on-chip bloco por bloco na SRAM durante o backward pass.
• Fusão de Kernels (Kernel Fusion): Fundir multiplicações de matrizes, mascaramento, softmax e dropout em um único kernel CUDA para eliminar viagens redundantes de ida e volta à HBM.

Método (com as equações principais)

Dados $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$, o FlashAttention define os tamanhos de bloco $B_c=\lceil M/(4d)\rceil$ e $B_r=\min (\lceil M/(4d)\rceil ,d)$, onde $M$ é a capacidade da SRAM.

Rastreamento Incremental do Softmax

Para um vetor decomposto em subvetores $x^{(1)}$ e $x^{(2)}$, as estatísticas cumulativas são atualizadas via: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

Para o bloco $j$ de chaves/valores e o bloco $i$ de consultas:

1. Computar pontuações não normalizadas na SRAM: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. Computar estatísticas locais do bloco: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$ e ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. Atualizar os normalizadores globais cumulativos: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. Atualizar o bloco de saída cumulativo ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

Análise de Complexidade

• Complexidade de IO: A atenção padrão requer $\Theta (Nd+N^2)$ acessos à HBM. O FlashAttention requer $\Theta (N^2{d}^2{M}^{-1})$ acessos à HBM. Os autores provam que nenhum algoritmo de atenção exata pode atingir assintoticamente $o(N^2{d}^2{M}^{-1})$ acessos para todos os valores de $M\in [d,Nd]$.
• FLOPs e Memória: O FlashAttention realiza $O(N^{2}d)$ FLOPs (com ligeiramente mais FLOPs no backward pass devido à recomputação) e requer apenas $O(N)$ de memória auxiliar além das entradas e saídas.

Extensão: Block-Sparse FlashAttention

Restringir a avaliação a blocos não nulos de acordo com uma máscara de esparsidade predefinida $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$ alcança uma complexidade de IO de $\Theta (Nd+N^2{d}^2{M}^{-1}s)$, onde $s$ é a fração de blocos não nulos.

Resultados (com os números mais importantes)

• Aceleração do Treinamento:
  • BERT-large (tam. seq. 512): Supera o recorde de velocidade de treinamento do MLPerf 1.1 em 15% (17,4 min vs. 20,0 min em $8\times$ A100s).
  • GPT-2 (tam. seq. 1K): Alcança aceleração de até 3,5x sobre o HuggingFace e até 1,8x sobre o Megatron-LM.
  • Long-Range Arena (LRA): Proporciona aceleração de 2,4x sobre a atenção padrão; o Block-sparse FlashAttention atinge 2,8x de aceleração.
• Capacidades com Contextos Mais Longos:
  • GPT-2 com contexto de 4K: Treina mais rápido que o Megatron-LM com contexto de 1K, enquanto melhora a perplexidade de validação em 0,7 pontos (17,5 vs. 18,2).
  • Classificação de Documentos Longos: Escalar o contexto de 512 para 16K resulta em um ganho de 4,3 no micro-$F_1$ no MIMIC-III (52,8 para 57,1) e de até 8,5 no ECtHR (72,2 para 80,7 com 8K).
  • Path-X e Path-256: O FlashAttention é o primeiro Transformer a resolver o Path-X (tam. seq. 16K) com 61,4% de acurácia. O Block-sparse FlashAttention alcança 63,1% de acurácia no Path-256 (tam. seq. 64K).

Limitações e questões em aberto

• Sobrecarga de Engenharia em CUDA: A implementação de kernels cientes de IO requer programação customizada de baixo nível em CUDA, em vez de abstrações padrão de PyTorch ou compiladores, tornando a portabilidade entre arquiteturas de hardware um desafio.
• Atenção Multi-GPU: O cronograma ótimo de IO em um único nó não se estende automaticamente para a atenção distribuída multi-GPU através de interconexões (NVLink, InfiniBand).
• Limite Inferior para Atenção Exata: Embora seja ótimo para $M$ de maneira geral, provar limites inferiores refinados de complexidade parametrizada em termos de $M$ continua sendo um problema em aberto.