Atenção exata, acelerada por nunca escrever a grade gigante. · Source

A mesma conta da atenção — só que nunca anotada.

O verdadeiro custo da atenção não é multiplicar — é arrastar uma tabela gigante de pontuações para dentro e para fora da memória lenta. Flash attention não muda nada nas contas e muda tudo sobre onde elas acontecem. A mesma resposta, numa fração do tempo. Como resolver um cubo em alta velocidade: o mesmo quebra-cabeça, os mesmos movimentos — a vantagem é puro método.

O gargalo não são as contas. É o transporte.

A atenção padrão monta a tabela completa de cada palavra pontuada contra todas as outras — depois escreve essa grade gigante na memória lenta, lê de volta para normalizar e lê de novo para ponderar os valores. Três idas e voltas com o maior objeto do modelo, enquanto o chip termina as contas e fica só esperando. Como um tapete desenrolado: grande demais para a sala, então fica no corredor — e cada olhada significa outra viagem.

Um chip veloz, sufocado por um cano lento.

T \approx \max\!\left(\dfrac{\text{FLOPs}}{\text{compute rate}},\ \dfrac{\text{bytes moved}}{\text{bandwidth}}\right)

Um processador faz aritmética muito mais rápido do que consegue buscar os números para mastigar. Então o tempo total é o que for mais lento: fazer as contas, ou mover os bytes. A atenção é limitada pela memória — as contas estão prontas, os dados não — então cortar viagens, e não cortar contas, é o que a acelera. Como uma bomba potente numa mangueira fina: poderia mover um oceano, mas é o cano que decide o fluxo.

Então não carregue tudo. Trabalhe um bloco de cada vez.

S_{ij} = \dfrac{Q_i K_j^{\top}}{\sqrt{d}}, \qquad Q_i,\, K_j \in \mathbb{R}^{B\times d}

A solução começa simples: nunca traga a grade inteira para perto. Fatie as Keys e os Values em blocos pequenos o bastante para caber no minúsculo bloco rápido do chip. Traga um bloco, faça a parte dele do trabalho, solte, busque o próximo. Em palavras simples: cada bloco de B linhas cabe na memória rápida, então você lida com apenas um par de blocos por vez. Como carregar lenha: você não arrasta a pilha inteira até a lareira — traz uma braçada.

Mas o softmax precisa da linha inteira de uma vez.

\mathrm{softmax}(x)_i = \dfrac{e^{\,x_i - m}}{\sum_{j=1}^{n} e^{\,x_j - m}}, \qquad m = \max_{j} x_j

Aqui está o problema. Para transformar pontuações em pesos, o softmax divide cada uma pelo total sobre cada key — e, por segurança, primeiro subtrai a maior pontuação da linha, para que as exponenciais nunca disparem. Tanto o máximo quanto a soma precisam da linha inteira. Mas você só tem um bloco. Como normalizar contra números que ainda não viu? Como servir porções justas de uma só jarra: você não pode encher o primeiro copo até saber quantos há.

O truque: manter um máximo corrente e reajustar conforme avança.

\begin{aligned} m_{\text{new}} &= \max(m_{\text{old}},\, \tilde m) \\ \ell_{\text{new}} &= e^{\,m_{\text{old}} - m_{\text{new}}}\,\ell_{\text{old}} + \sum_{i} e^{\,x_i - m_{\text{new}}} \end{aligned}

Você não precisa da linha inteira primeiro — precisa de uma contagem corrente. Acompanhe a maior pontuação vista até agora e o total em curso. Quando um novo bloco traz uma pontuação maior, ponha o total antigo na nova escala com um único fator de redução, e então incorpore o bloco; a saída se reajusta do mesmo jeito. O resultado é idêntico a fazer tudo de uma vez. Como a barra do salto em altura: erga-a quando chega um saltador melhor, e remeça cada marca anterior contra a nova altura.

A mesma resposta. A grade gigante nunca existiu.

\text{memory: } O(n^2)\ \longrightarrow\ O(n), \qquad \mathbf{o}_{\text{flash}} = \mathbf{o}_{\text{exact}}

Junte tudo e a tabela n×n nunca é armazenada — apenas um máximo corrente, uma soma corrente e a saída, todos pequenos. A memória cai de quadrática para linear, e a aceleração vem de muito menos viagens à memória lenta, não de cortar contas. E isto não é uma aproximação: o resultado é atenção exata. Como dobrar um tsuru de origami: uma folha, a forma pronta, nenhum retalho sobre a mesa.

🌱 Talvez nunca tenha sido o pensar que era lento.

Flash attention não acrescenta nenhuma ideia nova — apenas respeita onde a conta mora, reorganizando os mesmos pensamentos para o corpo que os abriga. Então, quanto da velocidade de uma mente é o pensar em si — e quanto é apenas o quão rápido ela alcança o que já sabe?