Resumo Técnico: SEKV – Cache de KV Adaptável à Resolução com Memória Semântica Hierárquica

Declaração do Problema

Grandes Modelos de Linguagem (LLMs) são cada vez mais implantados em cenários que exigem raciocínio de contexto longo (ex: análise de documentos, diálogo de múltiplos turnos, planejamento de agentes). No entanto, o cache de Chave-Valor (KV), que armazena estados intermediários para acelerar a decodificação, cresce linearmente com o comprimento da sequência. Reter um cache KV completo para contextos que excedem 128K tokens torna-se proibitivamente caro em termos de memória GPU, muitas vezes tornando a inferência de contexto longo inviável sem compressão.

As estratégias de compressão existentes enfrentam um compromisso fundamental entre eficiência e fidelidade:

1. Evicção de Tokens: Métodos que descartam permanentemente tokens com base em estatísticas de atenção (ex: H2O, SnapKV) correm o risco de fragmentar evidências coerentes ou perder irremediavelmente informações que se tornam críticas apenas mais tarde no processo de geração.
2. Agrupamento Semântico Estático: Métodos que comprimem tokens em blocos semânticos ou sentenças no momento do prefill (ex: ChunkKV, SentenceKV) fixam a resolução da compressão precocemente. Eles carecem da capacidade de recuperar detalhes de nível de token de um intervalo comprimido uma vez que esse intervalo se torne relevante durante a decodificação.

O desafio central é projetar uma organização de memória que preserve o acesso a evidências relevantes sob restrições severas de memória sem perda irreversível de informação, permitindo que o modelo ajuste dinamicamente a resolução do contexto armazenado com base na relevância da consulta.

Metodologia: SEKV

Os autores propõem o SEKV (Semantic KV Cache), um cache de KV semântico adaptável à resolução que organiza o contexto em uma memória semântica hierárquica abrangendo GPU e CPU. O SEKV evita a evicção irreversível de tokens armazenando o contexto em resoluções variáveis e "dando zoom" dinamicamente para o nível apropriado de detalhe durante a decodificação.

1. Segmentação de Intervalos Guiada por Entropia

Em vez de fragmentação de tamanho fixo, o SEKV segmenta o contexto de entrada em intervalos semânticos coerentes usando a surpresa do token (token surprisal) como um sinal de fronteira de custo zero.

• Mecanismo: Durante a passagem forward do prefill, a surpresa $H_t = -\log p(x_t | x_{<t})$ é computada para cada token. Tokens de alta surpresa, que frequentemente marcam transições semânticas (mudanças de tópico, introdução de entidades), são identificados como fronteiras de intervalo.
• Tokens Âncora: Esses tokens de fronteira de alta surpresa são mantidos em resolução total na GPU como "tokens âncora". Eles carregam informação semântica concentrada, essencial para roteamento e reconstrução posteriores.
• Intervalos (Spans): Os tokens entre dois âncoras vizinhos formam um intervalo semântico $S$.

2. Representação de Intervalo de Dupla Resolução e Hierarquia de Memória

Cada intervalo semântico é representado simultaneamente em duas resoluções através de uma hierarquia GPU–CPU:

• GPU (Roteamento Coarse): Cada intervalo é representado por um vetor de resumo leve (uma média ponderada pela surpresa e normalizada de vetores de chave projetados). Este resumo reside na GPU e é usado para um roteamento eficiente e de granularidade grossa para determinar quais intervalos são relevantes para a consulta atual.
• CPU (Reconstrução de Grão Fino): A estrutura completa de nível de token do intervalo é codificada de forma compacta na CPU usando uma base SVD de baixo posto. Um preditor leve atribui portões suaves (soft gates) por componente aos valores singulares, permitindo que o sistema adapte o orçamento de posto $r_S$ para cada intervalo com base em sua densidade de informação. Isso permite a reconstrução fiel, sob demanda, de pares KV de nível de token sem armazenar a matriz completa na GPU.

3. Mecanismo de Zoom-In Treinado

Durante a decodificação, o SEKV emprega um mecanismo aprendido para selecionar dinamicamente quais intervalos devem ser expandidos:

• Roteamento: A consulta atual é projetada no mesmo espaço de roteamento que os resumos dos intervalos. Um portão de relevância sigmoide calcula a probabilidade de um intervalo ser relevante.
• Expansão: Um limiar aprendido, por cabeça e por camada, determina se um intervalo será expandido. Se selecionado, o sistema busca assincronamente a base SVD de baixo posto correspondente da CPU, reconstrói os pares KV de nível de token e os integra ao cálculo de atenção.
• Atenção de Resolução Mista: O mecanismo de atenção opera em uma memória de resolução mista: tokens âncora contribuem com pares KV completos, intervalos não expandidos contribuem com seus vetores de resumo grosseiros (com uma correção de tamanho), e intervalos expandidos contribuem com pares de nível de token reconstruídos.
• Treinamento: O LLM base permanece congelado. Apenas as projeções de roteamento leves, os limiares de zoom-in e o preditor de portão de posto são treinados usando uma combinação de perda de destilação (alinhando com um professor de KV completo), supervisão de zoom, perda de reconstrução e uma penalidade de orçamento.

Principais Contribuições

1. Arquitetura SEKV: Um cache de KV semântico adaptável à resolução que armazena o contexto em múltiplas resoluções e ajusta dinamicamente o detalhe do cache com base na relevância da consulta, evitando a evicção irreversível de tokens.
2. Segmentação Guiada por Entropia: Um método que utiliza a surpresa do token para identificar intervalos semânticos coerentes e preservar tokens âncora de alta surpresa em resolução total para recuperação confiável.
3. Estrutura de Dupla Representação: Uma estrutura de KV inovadora que combina resumos de intervalos residentes na GPU para roteamento grosseiro, bases de baixo posto residentes na CPU para reconstrução aproximada de nível de token e um mecanismo de zoom-in treinado que expande seletivamente apenas os intervalos relevantes à consulta.

Resultados Experimentais

O SEKV foi avaliado em quatro benchmarks (LongBench, RULER, InfiniteBench, Needle-in-a-Haystack) e cinco modelos de base (incluindo variantes Llama-3, Mistral e Qwen).

• Desempenho: O SEKV supera o baseline de compressão semântica mais forte (SentenceKV) em uma média de 5,9% em todos os benchmarks e modelos. Ele alcança os melhores resultados em 20 de 20 comparações benchmark-modelo.
• Fidelidade de Recuperação: No benchmark Needle-in-a-Haystack (NIAH), o SEKV mostra uma estabilidade de recuperação superior através de variados comprimentos de contexto e profundidades de agulha em comparação com métodos de evicção de tokens e métodos semânticos estáticos.
• Eficiência de Memória: Em um contexto de 128K, o SEKV reduz o consumo de memória GPU em 53,3% em comparação com o cache KV completo, mantendo o desempenho próximo ao baseline não comprimido.
• Eficiência: O SEKV alcança latência e throughput competitivos, muitas vezes superando o cache KV completo em latência devido à redução da pressão de largura de banda de memória, enquanto supera significativamente os métodos de evicção agressiva em precisão.
• Eficiência de Parâmetros: O método adiciona menos de 0,05% de parâmetros treináveis (aprox. 4,3 milhões de parâmetros para Llama-3-8B) mantendo o modelo base congelado.

Significância e Alegações

O artigo argumenta que a inferência eficiente de contexto longo requer uma mudança da compressão estática para a recuperação dinâmica de resolução. Os autores afirmam que, enquanto a compressão estática preserva o contexto grosseiro, ela falha em recuperar evidências de grão fino quando estas se tornam relevantes tardiamente no processo de geração. O SEKV aborda isso organizando o contexto em intervalos semânticos guiados por entropia e utilizando um sistema de memória hierárquica que permite ao modelo "dar zoom" na informação relevante sob demanda.

A significância deste trabalho reside em demonstrar que:

1. A evicção irreversível não é necessária: A inferência de contexto longo de alta fidelidade pode ser alcançada ao descarregar representações de baixo posto recuperáveis para a CPU em vez de descartar tokens.
2. A resolução adaptável é crítica: Alocar resolução de grão fino apenas onde necessário (via mecanismo de zoom-in) permite que os modelos equilibrem restrições de memória com fidelidade de recuperação de forma mais eficaz do que políticas estáticas.
3. A estrutura semântica auxilia a compressão: O uso de segmentação guiada por entropia cria unidades mais compressíveis do que blocos de tokens arbitrários, permitindo uma aproximação de baixo posto mais eficiente.

Os autores reconhecem limitações, observando que a eficácia depende da qualidade da construção dos intervalos semânticos (que pode ter dificuldades com discursos altamente fragmentados) e que a abordagem introduz uma dependência da largura de banda CPU-GPU, o que pode impactar a latência em cenários de consulta adversos. No entanto, eles concluem que o SEKV fornece um framework robusto para a inferência de LLM de contexto longo escalável e de alta fidelidade.