o problema
Attention custa tempo e memória quadráticos no comprimento da sequência. A resposta da comunidade entre 2019 e 2022 foi aproximar: esparsidade, baixo posto, hashing, combinações disso. Quase todos esses métodos entregam complexidade linear ou quase linear em FLOP e, mesmo assim, quase nenhum ficou mais rápido no relógio. Reduzir conta não adiantou porque conta não era o gargalo.
Numa A100, a HBM tem 40 a 80 GB a 1,5-2,0 TB/s; a SRAM on-chip tem 192 KB por multiprocessador, 108 deles, a algo em torno de 19 TB/s. Uma ordem de grandeza mais rápida, muitas ordens de grandeza menor. A implementação padrão de attention calcula S = QKᵀ, escreve uma matriz N×N na HBM, lê de volta para a softmax, escreve P na HBM, lê de novo para multiplicar por V. Softmax, masking e dropout são operações memory-bound: o tempo delas é o tempo de ir e voltar da HBM. Para GPT-2 medium com sequência 1.024, isso dá 40,3 GB de tráfego contra 66,6 GFLOPs de conta. O framework em Python também não ajuda: PyTorch e TensorFlow não dão controle fino sobre acesso à memória.
a ideia
Tratar attention como um problema de I/O, não de aritmética. A literatura de banco de dados, álgebra linear numérica e processamento de imagem já contava leituras e escritas entre níveis de memória há décadas; deep learning não contava.
A meta passa a ser nunca escrever a matriz N×N na HBM. Isso exige duas coisas: calcular a softmax sem ver a linha inteira de uma vez, e não guardar a matriz intermediária para o backward. As duas têm solução conhecida — tiling e recomputação — e o trabalho aqui foi juntá-las num único kernel CUDA. O resultado é exato, bit a bit equivalente na prática: as curvas de perplexidade de validação do GPT-2 ficam em cima das do baseline.
como funciona
A softmax é decomponível se você carregar duas estatísticas por linha: o máximo corrente m e a soma corrente ℓ. Ao juntar dois blocos, reescala-se o que já foi acumulado pelo novo máximo e soma-se o novo pedaço.
O laço externo percorre blocos de K e V; o interno, blocos de Q. Os blocos são dimensionados para caber na SRAM: Bc = ⌈M/4d⌉ e Br = min(⌈M/4d⌉, d), com M o tamanho da SRAM.
for j in blocos_de_KV:
K_j, V_j = load(K[j]), load(V[j]) # HBM -> SRAM, uma vez
for i in blocos_de_Q:
Q_i, O_i, l_i, m_i = load(...)
S_ij = tau * Q_i @ K_j.T # Br x Bc, nasce e morre na SRAM
m_til = rowmax(S_ij)
P_ij = exp(S_ij - m_til)
m_new = max(m_i, m_til)
l_new = exp(m_i - m_new)*l_i + exp(m_til - m_new)*rowsum(P_ij)
O_i = (l_i*exp(m_i - m_new)*O_i + exp(m_til - m_new)*(P_ij @ V_j)) / l_new
store(O_i, l_new, m_new)O backward guarda apenas O, as estatísticas (m, ℓ) e o estado do gerador pseudoaleatório do dropout, e reconstrói S e P na SRAM. Um detalhe salva a conta do gradiente da softmax: Dᵢ = Pᵢ·dPᵢ é igual a doᵢ·oᵢ, um produto interno de vetores de tamanho d em vez de uma redução sobre N.
A conta de I/O: Θ(Nd + N²) acessos à HBM na versão padrão contra Θ(N²d²/M) aqui. Com d entre 64 e 128 e M na casa de 100 KB, d² é muitas vezes menor que M — daí os até 9 vezes menos acessos. A Proposição 3 fecha a porta: não existe algoritmo de attention exata com o(N²d²/M) acessos para todo M em [d, Nd].
Os resultados: BERT-large 15% mais rápido que o recorde do MLPerf 1.1 (17,4 contra 20,0 minutos), GPT-2 small em 2,7 dias contra 9,5 do HuggingFace e 4,7 do Megatron-LM, LRA 2,4× mais rápido. Com contexto 4K, o GPT-2 treina 30% mais rápido que o Megatron com contexto 1K e ganha 0,7 de perplexidade. Path-X (16K) sai de aleatório para 61,4%.
o que isso custou
Cada variante de attention vira um kernel CUDA escrito à mão. É engenharia cara, num nível bem abaixo do PyTorch, e a implementação não migra sozinha entre arquiteturas de GPU — os autores pedem explicitamente um compilador que faça isso a partir de linguagem alta, tipo Halide.
A recomputação cobra FLOPs. Contra o Apex FMHA, o backward do FlashAttention é mais lento (0,20 contra 0,17 ms em sequência 128); o ganho vem do forward e do fim da materialização.
O ganho depende do hardware. Em T4, com SRAM menor, os blocos encolhem e a aceleração cai. Com head dimension 128 também cai, pelo mesmo motivo.
E continua quadrático em FLOP. A partir de sequências entre 512 e 1.024, métodos aproximados como o Linformer começam a ficar mais rápidos. A resposta do paper é o block-sparse FlashAttention — que é aproximação, com o custo de qualidade correspondente (59,6 contra 59,8 de média no LRA).
Por fim, a otimalidade vale para uma GPU. Atenção distribuída entre GPUs adiciona um nível de hierarquia que este trabalho não analisa.
onde isso aparece hoje
O código saiu aberto em HazyResearch/flash-attention e virou peça de infraestrutura: hoje o PyTorch expõe scaled_dot_product_attention com backend de flash attention, e o algoritmo ganhou versões sucessoras afinadas para gerações novas de GPU. A vocabulário mudou junto: “IO-aware” e “quantos bytes isso move” entraram na conversa sobre kernel, ao lado de “quantos FLOPs isso faz”.
O efeito mais visível foi no contexto. O paper mostra o mecanismo direto — memória linear na sequência permite treinar com janela maior pelo mesmo orçamento, e janela maior melhora o modelo: 4,3 pontos de F1 em MIMIC-III indo de 512 para 16K tokens, 8,5 pontos em ECtHR indo de 512 para 8K. A corrida por janelas grandes que veio depois começa por remover essa barreira de memória.