Destrinchando os LLMs (post 4): KV Cache, o motor invisível da inferência.

Mais um post da série em que destrincho as entranhas dos LLMs.

Nos textos anteriores conversamos sobre embeddings, mecanismo de atenção e janela de contexto. Hoje quero abordar um componente muito menos conhecido, mas absolutamente fundamental para que os LLMs sejam utilizáveis na prática: o KV Cache (Key-Value Cache).

Se você já se perguntou por que um LLM normalmente demora alguns segundos para começar a responder, mas depois passa a gerar texto de forma relativamente contínua, boa parte da explicação está nesse mecanismo. Embora praticamente invisível para o usuário, o KV Cache é um dos principais responsáveis por tornar a inferência suficientemente rápida e economicamente viável.

Para entender sua função, vale lembrar como um LLM autorregressivo baseado em um Transformer decoder-only, como GPT, Claude ou Llama, produz uma resposta.

A geração ocorre em duas etapas bastante diferentes. A primeira é conhecida como prefill. Nessa fase, o modelo processa todo o contexto disponível, como prompt, instruções do sistema, histórico da conversa, documentos e demais informações enviadas na requisição. Como todos esses tokens já são conhecidos, suas representações podem ser computadas de forma amplamente paralela, respeitadas as limitações do hardware e da implementação utilizada.

A segunda etapa é a geração propriamente dita. Ela ocorre de forma autorregressiva: o modelo produz um token por vez. Cada token recém-gerado passa imediatamente a integrar o contexto utilizado para gerar o próximo. Essa dependência faz com que essa fase seja inerentemente sequencial.

É justamente aqui que surge um enorme desafio computacional. Sempre que um novo token é produzido, o mecanismo de atenção precisa considerar todo o contexto disponível naquele instante. Se o modelo tivesse de reconstruir todas as representações internas dos milhares de tokens anteriores a cada novo passo, grande parte do trabalho seria repetida continuamente.

É exatamente esse desperdício que o KV Cache elimina. Durante o processamento do contexto e da resposta em construção, cada camada do Transformer projeta as representações dos tokens em três conjuntos de vetores conhecidos como Queries (Q), Keys (K) e Values (V). O mecanismo de atenção utiliza essas representações para determinar quais informações presentes no contexto devem influenciar a geração do próximo token.

Nos LLMs autorregressivos baseados em atenção causal, existe uma propriedade importante: os Keys e Values produzidos para um determinado token dependem apenas desse token e dos tokens anteriores, nunca daqueles que ainda serão gerados. Em outras palavras, o futuro não altera retrospectivamente essas representações.

É justamente essa característica da atenção causal que torna o KV Cache possível. Como os Keys e Values dos tokens anteriores permanecem válidos durante toda a inferência, eles podem ser armazenados temporariamente e reutilizados nas etapas seguintes, em vez de serem recalculados repetidamente.

Quando chega o momento de produzir um novo token, o modelo continua atravessando todas as suas camadas para calcular as representações correspondentes ao token recém-adicionado, incluindo sua Query, Key e Value. Entretanto, não precisa reconstruir os Keys e Values de todos os tokens anteriores. Esses já estão disponíveis no KV Cache. O ganho de desempenho vem exatamente dessa reutilização.

É importante observar que o KV Cache não elimina o mecanismo de atenção. A cada novo token, a Query correspondente ao token recém-gerado ainda precisa ser comparada com todos os Keys armazenados no contexto para determinar quais informações deverão ser recuperadas a partir dos respectivos Values. O cache elimina o recálculo das representações anteriores, mas não elimina as comparações inerentes ao mecanismo de atenção.

Essa distinção costuma gerar confusão. O KV Cache não torna a geração paralela. Ela continua sendo sequencial, token por token. Tampouco aumenta a janela de contexto, cria memória permanente ou faz o modelo “lembrar” de conversas anteriores. Trata-se exclusivamente de uma estratégia de reutilização de cálculos durante uma única inferência.

Esse detalhe aparentemente técnico ajuda a explicar vários comportamentos observados no uso cotidiano dos LLMs. Por exemplo, normalmente existe um intervalo relativamente maior entre o envio do prompt e o início da resposta. Esse tempo corresponde, em grande parte, ao processamento inicial do contexto, o prefill. Depois que esse trabalho foi concluído e o KV Cache foi construído, boa parte dos cálculos deixa de ser repetida, permitindo que os tokens seguintes sejam produzidos muito mais rapidamente.

O KV Cache também ajuda a entender um aspecto importante da economia da IA generativa. Para um modelo fixo, o espaço ocupado pelo KV Cache cresce aproximadamente de forma linear com o número de tokens presentes no contexto, multiplicado pelo número de camadas e pelas dimensões internas da arquitetura. Isso significa que contextos maiores não consomem apenas mais processamento durante o prefill; eles também exigem mais memória para manter armazenadas as representações internas que permanecerão disponíveis durante toda a geração.

Nos modelos mais modernos, com dezenas de bilhões de parâmetros e janelas de contexto extremamente extensas, o KV Cache pode consumir dezenas ou até centenas de gigabytes de memória, tornando-se um dos principais fatores de custo da inferência em GPUs.

Existe, entretanto, um limite importante. Embora o KV Cache elimine praticamente todo o recálculo redundante dos tokens anteriores, cada novo token ainda precisa comparar sua Query com um conjunto crescente de Keys armazenados. Assim, à medida que o contexto aumenta, a etapa dominante deixa de ser reconstruir o passado e passa a ser consultar um número cada vez maior de representações armazenadas. Como consequência, a latência por token tende a crescer gradualmente em contextos muito longos, mesmo utilizando cache.

Em muitos ambientes de produção, a quantidade de memória disponível nas GPUs torna-se um dos principais fatores limitantes da inferência. Não basta apenas possuir capacidade de processamento. É necessário manter o KV Cache residente na memória durante toda a geração da resposta. Por isso, ampliar continuamente a janela de contexto implica custos crescentes de infraestrutura.

Esse é um dos motivos pelos quais empresas e grupos de pesquisa investem intensamente em técnicas de compressão, quantização, descarte seletivo (eviction), compartilhamento e gerenciamento eficiente do KV Cache. Pequenos ganhos de eficiência podem representar reduções expressivas no custo operacional quando milhões de inferências são executadas diariamente.

Vale também não confundir KV Cache com Prompt Caching. Embora ambos procurem reduzir custos, tratam problemas diferentes. O KV Cache reutiliza representações internas calculadas durante uma única inferência. Já o Prompt Caching, implementado de maneiras distintas pelos diferentes provedores, procura evitar o reprocessamento de partes idênticas de prompts entre inferências distintas quando determinadas condições são satisfeitas. São mecanismos complementares, mas conceitualmente diferentes.

O usuário raramente percebe que tudo isso está acontecendo. Quando um modelo responde de forma fluida, tendemos a atribuir essa velocidade apenas à potência das GPUs. Na realidade, ela resulta da combinação de diversos mecanismos de engenharia, dos quais o KV Cache é um dos mais importantes.

No fundo, o KV Cache ilustra um princípio clássico da ciência da computação: muitas vezes, o maior ganho de desempenho não vem de realizar cálculos mais rapidamente, mas de evitar recalcular aquilo que já foi computado. É um detalhe de implementação quase invisível para quem utiliza um LLM, mas essencial para tornar a inferência economicamente viável. Sem ele, responder token por token exigiria repetir uma enorme quantidade de trabalho já realizado, tornando os modelos significativamente mais lentos, mais caros e muito menos escaláveis.