KV Cache

O cache KV (cache de chave-valor) é uma técnica de otimização essencial, usada principalmente em grandes modelos de linguagem (LLMs) e outras arquiteturas baseadas em Transformer para acelerar a latência de inferência e reduzir os custos computacionais. Em sua essência, o cache KV armazena as matrizes Key e Value geradas pelo mecanismo de atenção para tokens anteriores de uma sequência. Ao salvar esses cálculos intermediários, o modelo evita recalcular os estados de atenção de todo o histórico da conversa cada vez que gera um novo token. Esse processo transforma o fluxo de trabalho de geração de texto de uma operação de complexidade quadrática em uma operação linear, viabilizando interações em tempo real com chatbots e agentes de IA.

Mecanismo e benefícios#

Em um modelo Transformer padrão, para gerar a próxima palavra é preciso considerar todas as palavras anteriores a fim de entender o contexto. Sem o cache, o modelo precisaria recalcular as relações matemáticas de toda a sequência a cada etapa. O cache KV resolve isso funcionando como um banco de memória.

• Aumento de velocidade: Ao recuperar da memória chaves e valores pré-computados, o sistema acelera drasticamente o mecanismo de inferência. Isso é essencial para aplicações que exigem baixa latência, como a inferência em tempo real em bots de atendimento ao cliente.
• Eficiência de recursos: Embora aumente o uso de memória (VRAM), o cache reduz significativamente a computação (FLOPs) necessária por token. Esse equilíbrio costuma ser gerenciado com técnicas como a quantização de modelos ou a paginação, de forma semelhante ao gerenciamento de RAM pelos sistemas operacionais.
• Contexto ampliado: O gerenciamento eficiente do cache KV permite que os modelos lidem com uma janela de contexto maior, possibilitando o processamento de documentos longos ou a manutenção de conversas coerentes por períodos prolongados.

Aplicações no mundo real#

O cache KV é um componente fundamental na implantação de IA generativa moderna, mas seus princípios também se aplicam à visão computacional (CV).

1. Chatbots generativos: Serviços como o ChatGPT ou o Claude dependem muito do cache KV. Quando um usuário faz uma pergunta complementar, o modelo não relê todo o histórico da conversa desde o início. Em vez disso, ele acrescenta a nova entrada aos estados em cache da interação anterior, permitindo respostas quase instantâneas.

2. Compreensão de vídeo: Em tarefas de compreensão de vídeo, os modelos processam os quadros em sequência. Assim como os tokens de texto, as características visuais de quadros anteriores podem ser armazenadas em cache para ajudar o modelo a rastrear objetos ou reconhecer ações sem reprocessar todo o histórico do vídeo. Isso é especialmente relevante para o reconhecimento de ações, em que o contexto temporal é crucial.

Gerenciamento eficiente de memória#

À medida que os modelos crescem, o tamanho do cache KV pode se tornar um gargalo, consumindo gigabytes de memória da GPU. Os avanços recentes se concentram em otimizar esse armazenamento.

• PagedAttention: Inspirado na memória virtual dos sistemas operacionais, o PagedAttention (apresentado pelo vLLM) permite armazenar o cache KV em blocos de memória não contíguos. Isso reduz a fragmentação e permite tamanhos de lote maiores durante o serving de modelos.
• Quantização do cache KV: Para economizar espaço, os desenvolvedores costumam aplicar precisão mista ou quantização int8 especificamente aos valores em cache. Isso reduz a memória ocupada, permitindo que dispositivos de IA de borda com RAM limitada executem modelos capazes.
• Cache de prompt: Técnica relacionada em que os estados KV de um prompt de sistema estático (por exemplo, "Você é um assistente de programação prestativo") são calculados uma vez e reutilizados em várias sessões de usuário. Esse é um recurso fundamental para otimizar fluxos de trabalho de engenharia de prompts em grande escala.

Distinguir conceitos relacionados#

É útil diferenciar o cache KV de outros termos relacionados a cache e otimização:

• Cache KV vs. cache de prompt: O cache KV normalmente se refere à memória dinâmica, token a token, usada durante um único fluxo de geração. O cache de prompt se refere especificamente ao armazenamento do estado processado de uma instrução de entrada fixa, para reutilização em várias chamadas de inferência independentes.
• Cache KV vs. incorporações: As incorporações são representações vetoriais de dados de entrada (texto ou imagens) que capturam o significado semântico. O cache KV armazena as ativações (chaves e valores) derivadas dessas incorporações nas camadas de atenção, especificamente para a geração de sequências.
• Cache KV vs. pesos do modelo: Os pesos do modelo são os parâmetros estáticos e aprendidos da rede neural. O cache KV consiste em dados dinâmicos e temporários gerados durante a passagem direta de uma sequência de entrada específica.

Exemplo: contexto em modelos de visão#

Embora o cache KV seja mais conhecido em PLN, o conceito de manutenção de estado também se aplica a modelos avançados de visão. No exemplo abaixo, simulamos a ideia de transmitir estado (contexto) em um cenário de rastreamento de vídeo usando o Ultralytics YOLO26. Aqui, o rastreador mantém a identidade dos objetos entre os quadros, de forma conceitualmente semelhante a como um cache mantém o contexto entre tokens.

from ultralytics import YOLO

# Load the Ultralytics YOLO26 model
model = YOLO("yolo26n.pt")

# Track objects in a video, maintaining identity state across frames
# The 'track' mode effectively caches object features to link detections
results = model.track(source="https://ultralytics.com/images/bus.jpg", show=False)

# Print the ID of the tracked objects
if results[0].boxes.id is not None:
    print(f"Tracked IDs: {results[0].boxes.id.numpy()}")

Os desenvolvedores que querem gerenciar conjuntos de dados e implantar modelos otimizados podem usar a Ultralytics Platform, que simplifica o fluxo de trabalho desde a anotação de dados até a implantação eficiente de modelos. Para quem tem interesse nos mecanismos mais profundos da atenção, bibliotecas como o PyTorch fornecem os blocos fundamentais em que esses mecanismos de cache são implementados.