Última atualização: 2026-06-22

Memória da atenção

A atenção padrão de transformers compara cada token com todos os outros tokens, construindo uma matriz de scores antes do softmax. Essa matriz é quadrada no comprimento da sequência, então sua memória cresce com o quadrado do contexto — o famoso custo quadrático da atenção. Esta calculadora estima o tamanho dessa matriz materializada a partir do comprimento da sequência, do número de cabeças de atenção, do tamanho do lote e dos bytes usados por score. É precisamente a memória que o FlashAttention evita armazenar.

O termo quadrático

Para uma sequência de $s$ tokens, a atenção forma uma matriz de scores $s \times s$: a linha $i$, coluna $j$ guarda o quanto o token $i$ atende ao token $j$. Armazená-la exige memória proporcional a $s^2$. Isso difere dos pesos do modelo ou do cache de chave-valor, que crescem linearmente com o contexto — a matriz de scores cresce com o quadrado, então, em contextos longos, ela se torna o buffer intermediário dominante e a razão pela qual a atenção ingênua esgota a memória.

A fórmula

Cada cabeça constrói sua própria matriz de scores e cada sequência do lote carrega sua própria cópia, então a memória em bytes é

onde $B$ é o tamanho do lote, $H$ o número de cabeças de atenção, $s$ o comprimento da sequência e $e$ os bytes por score armazenado. Dividir por $10^9$ resulta em gigabytes. O quadrado em $s$ é o que torna o contexto longo custoso: o número de cabeças e o lote multiplicam apenas linearmente.

O que é essa memória

O valor mede a matriz de scores materializada que a atenção padrão escreve na memória entre o produto consulta-chave e a soma ponderada pelo softmax sobre os valores. É o buffer de ativação que uma implementação ingênua precisa manter. Se ele persiste depende do cenário: em uma passagem direta simples, é transitório e pode ser reutilizado entre as camadas, enquanto os frameworks de treinamento podem manter uma cópia por camada para a passagem reversa, a menos que o checkpointing ou um kernel fundido intervenha.

FlashAttention

O FlashAttention calcula a saída idêntica sem nunca escrever a matriz $s \times s$ completa. Ele percorre as chaves e os valores em pequenos blocos, mantendo estatísticas correntes do softmax, de modo que só precisa de um bloco na memória por vez, o que transforma o armazenamento de quadrático em linear no comprimento da sequência. A matriz cujo tamanho esta calculadora mede é exatamente o que o FlashAttention se recusa a armazenar — então o resultado aqui é um bom indicador da memória que um kernel de atenção fundido economiza em um dado comprimento de contexto.

Exemplo resolvido

Considere uma única sequência de 4.096 tokens com 32 cabeças de atenção em precisão de 16 bits:

cerca de 1,07 GB para a matriz de scores de uma camada. Dobre o contexto para 8.192 tokens e o termo quadrático assume o controle: a mesma expressão dá aproximadamente 4,29 GB, quatro vezes mais para o dobro do comprimento. O lado de chave-valor do mesmo orçamento de atenção é coberto na Calculadora de Tamanho do Cache KV, e a parte dos pesos do modelo na Calculadora de VRAM para Inferência de LLM.