Ir para o conteúdo

The Cost of Attention

O caminho do texto ao próximo token em sete passos: tokens, embeddings, query/key/value, máscara e cabeças, o bloco, a saída e a geração. ATENÇÃO TRANSFORMERS Tokenstexto → ids Embeddingsids → X Q · K · Vuma cabeça Máscaraduas cabeças BlocoLN · MLP Saídalogits → p Geraçãoescolher texto → → próximo token

Uma página à parte. O exemplo destas páginas tem cinco tokens e não custa nada para rodar. Esta página trata do que a atenção custa num modelo de verdade: o cache que a geração guarda de um token para o seguinte e o \(O(n^2)\) dos contextos longos. ← Máscaras e cabeças · Visão geral

O cache de KV é o que ninguém ensina e todo mundo paga

Um modelo gera texto um token por vez e cada token novo precisa de atenção sobre todos os anteriores. O que isso exige de cada token antigo? A query dele não: ela só serviu para calcular a própria saída, lá atrás. Mas a chave e o valor dele serão consultados por todo token que vier depois. Sem guardá-los, o passo \(t\) rodaria a rede inteira sobre os \(t\) tokens do prefixo só para refazê-los e gerar \(n\) tokens custaria \(1 + 2 + \dots + n = O(n^2)\) passagens de token em vez de \(n\) — então ninguém faz isso. Cada K e V é calculado uma vez e guardado.

Esse cache é hoje o custo dominante de memória para servir um modelo e é a razão pela qual "contexto de 128k" é uma decisão de preço, não de arquitetura:

\[ \text{bytes de cache} \;=\; 2 \times L \times h_{kv} \times d_{\text{head}} \times n \times B \times \text{bytes} \]

O 2 conta K e V; \(L\) é o número de camadas, \(n\) o comprimento do contexto e \(B\) o tamanho do batch. Ponha números de um modelo real, o Llama 3 8B: \(L = 32\), \(h_{kv} = 8\), \(d_\text{head} = 128\), em 16 bits (2 bytes). Cada token custa \(2 \times 32 \times 8 \times 128 \times 2 = 131\,072\) bytes, ou 128 KiB. Um contexto de 128k tokens ocupa então 16 GiB — tanto quanto os próprios pesos do modelo — e isso para um usuário. O brinquedo destas páginas, com uma camada e duas cabeças de largura 2, guarda \(2 \times 1 \times 2 \times 2 = 8\) números por token; o GPT-2 small guarda \(2 \times 12 \times 12 \times 64 = 18.432\), 36 KiB por token em 16 bits.

Só um termo desse produto é uma variável livre de projeto: \(h_{kv}\), o número de cabeças de chave/valor. As queries podem manter todas as \(h\) cabeças; K e V podem ser compartilhados. O Llama 3 8B tem 32 cabeças de query e só 8 de chave/valor — com uma por query o cache seria quatro vezes maior. Na tabela de formas de máscaras e cabeças, K e V passam a ser \(B \times h_{kv} \times n \times d_\text{head}\) e cada um é repetido \(h / h_{kv}\) vezes para encontrar as queries; só a versão pequena vai para o cache. Isso produziu a sequência sobre a qual todo modelo atual se posiciona:

Esquema Cabeças de KV Cache Qualidade Situação em 2026
MHA — multi-cabeça1 \(h\) 1× Referência Histórico. Ninguém serve isso em contexto longo.
MQA — multi-query3 1 \(1/h\) Mensuravelmente pior Usado onde memória domina tudo
GQA — grouped-query4 \(g\) (tipicamente 8) \(g/h\) Indistinguível do MHA O padrão. Llama, Mistral, Qwen, Gemma
MLA — latente multi-cabeça5 — ~\(1/57\) Reportada igual ou acima do MHA DeepSeek-V2/V3; comprime KV em um latente de posto baixo e faz cache dele

Isto é um problema de banda de memória, não de FLOPs

Durante a geração você lê o cache de KV inteiro para produzir um token. A intensidade aritmética é péssima, o acelerador fica ocioso e a vazão é ditada pela banda da HBM. Tudo que acelera inferência — GQA, cache de KV quantizado, PagedAttention6, decodificação especulativa, batches maiores — é um ataque a bytes movidos, não a operações realizadas. Reconhecer de qual recurso você está com falta é a maior parte da engenharia prática de LLMs.

Tornando o \(O(n^2)\) pagável

A atenção é \(O(n^2 d)\) em tempo e, escrita ingenuamente, \(O(n^2)\) em memória: a matriz \(S\) da tabela de formas tem uma entrada por par de tokens. Com \(n = 128\text{k}\) ela tem \(1{,}7 \times 10^{10}\) entradas — 34 GB em 16 bits, por cabeça, por camada. Existem quatro famílias de respostas e elas não têm o mesmo sucesso.

  • Mesma matemática, execução melhor

    O FlashAttention2 nunca materializa a matriz \(n \times n\). Ele fatia a computação em blocos que cabem na SRAM e usa o truque do softmax online para combiná-los. A memória cai para \(O(n)\), o tempo de parede cai várias vezes e a saída é a da versão ingênua a menos do arredondamento de ponto flutuante.

    Isto não é uma aproximação e não há trade-off de qualidade. O FlashAttention-2 reorganiza o trabalho na GPU e o FlashAttention-3 o afina para hardware classe Hopper. Se você está escrevendo atenção à mão, é quase certo que está mais lento — use F.scaled_dot_product_attention.

  • Atender a menos

    Janela deslizante: cada token vê apenas os últimos \(w\); empilhar camadas amplia o alcance efetivo exatamente como o campo receptivo de uma CNN. Sumidouros de atenção7: manter os primeiros tokens sempre visíveis — os modelos despejam massa de atenção excedente neles e removê-los destrói a geração.

    Modelos modernos de contexto longo intercalam: algumas camadas de atenção completa entre muitas locais, de modo que a informação global ainda tem rota enquanto o custo fica quase linear.

  • Mudar a matemática

    A atenção linear troca o softmax por um kernel, tornando a operação associativa e o custo \(O(n)\) — e abrindo mão da recuperação afiada que o softmax oferece. Modelos de espaço de estados (Mamba, Mamba-28) alcançam o mesmo \(O(n)\) com um estado recorrente e um cache de tamanho constante.

    O veredito honesto de 2026: modelos lineares puros perdem em recuperação exata de contexto longo. Híbridos vencem — Jamba, Zamba, Nemotron-H e Falcon-H1 combinam camadas SSM com uma parcela pequena de atenção completa (uma minoria das camadas, ou cabeças de atenção ao lado do SSM) e reportam qualidade de Transformer com uma fração do cache.

  • Aprender o que pular

    Atenção esparsa treinável — o NSA da DeepSeek9 e o MoBA da Kimi — escolhe quais blocos do passado atender, com a seleção treinada ponta a ponta em vez de fixada por um padrão.

    Esta é a mais promissora das quatro no momento: mantém a recuperação exata do softmax onde importa e paga custo quase linear no resto — e, ao contrário de esquemas esparsos anteriores, é alinhada ao hardware o bastante para de fato ser mais rápida.

O que usar, na prática

Atenção completa com FlashAttention e GQA, RoPE para posição, janela deslizante ou camadas híbridas se suas sequências forem genuinamente longas. Essa combinação é o que essencialmente todo modelo lançado nos últimos dois anos faz e sair dela exige uma razão medida.

Verifique você mesmo — enquanto o brinquedo gera ␣telhado depois de “o gato subiu no”, o que está no cache e o que é calculado?

As chaves e os valores dos cinco tokens do prompt, na sua única camada e nas duas cabeças: \(5 \times 8 = 40\) números. Para ␣telhado o modelo calcula só a linha nova (o embedding, \(q\), \(k\) e \(v\), uma linha de scores contra seis chaves e o MLP) e acrescenta \(k\) e \(v\) dela ao cache. Sem o cache ele recalcularia as seis linhas; gerar \(n\) tokens custaria \(1 + 2 + \dots + n\) linhas.

← Máscaras e cabeças · Visão geral · próximo capítulo: Transformers

  1. Vaswani, A., et al. (2017). Attention Is All You Need — NeurIPS. Atenção escalada por produto escalar, multi-cabeça, posições senoidais. ↩

  2. Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — NeurIPS. Veja também FlashAttention-2 (2023) e FlashAttention-3 (2024). ↩

  3. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. Atenção multi-query e a primeira afirmação clara de que decodificar é limitado por memória. ↩

  4. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — EMNLP. O meio-termo que virou padrão. ↩

  5. DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. Atenção latente multi-cabeça: fazer cache de um latente de posto baixo em vez de K e V. ↩

  6. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention — SOSP. Memória virtual para o cache de KV; o núcleo do vLLM. ↩

  7. Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024). Efficient Streaming Language Models with Attention Sinks — ICLR. Por que o primeiro token não pode ser descartado. ↩

  8. Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality — ICML. Mamba-2 e a equivalência que torna híbridos projetáveis. ↩

  9. Yuan, J., et al. (2025). Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention — ACL. Esparsidade aprendida durante o pré-treino em vez de imposta depois dele. ↩