Resumo Técnico: LOCKS – Sumários de Chaves Compactos por Página para Decodificação Eficiente de Contexto Longo

1. Declaração do Problema

Servir Modelos de Linguagem de Grande Escala (LLMs) em comprimentos de contexto longos é cada vez mais limitado pelo cache de Chave-Valor (KV), especificamente pela largura de banda de leitura de KV, e não pela capacidade de memória. Em regimes de produção com alta concorrência, os pesos do modelo são amortizados em um lote (batch), mas cada requisição deve ler seu próprio cache KV único. Para modelos como o Llama-3.1-8B, um contexto de 128K requer ~16 GiB de cache KV por requisição, excedendo o tamanho dos pesos do modelo. Consequentemente, o cache, e não os pesos, limita o número de requisições que um acelerador pode servir.

Embora a atenção seja globalmente de alto posto (high-rank), ela é localmente de baixo posto (low-rank): algumas páginas (blocos de tokens) carregam quase toda a massa de atenção de uma consulta (query) em qualquer etapa de decodificação. A decodificação esparsa visa atender apenas a essas páginas críticas. No entanto, os métodos de seleção existentes enfrentam um compromisso:

• Seleção exata (ler todas as chaves) anula o propósito da esparsidade ao incorrer nos custos totais de largura de banda.
• Sumários residentes (representações compactas de páginas) são necessários para evitar a leitura de chaves completas, mas as abordagens existentes frequentemente falham em preservar a estrutura geométrica específica necessária para identificar páginas "portadoras" (aquelas que contêm a resposta) sob restrições de orçamento agressivas. Projeções compartilhadas entre páginas são comprovadamente cegas ao conteúdo específico da página, levando ao colapso de tarefas em cenários de recuperação e raciocínio.

2. Metodologia: LOCKS

LOCKS (Page-Local Compact Key Summaries) aborda isso mantendo o cache KV completo residente, mas evitando a leitura da maioria das chaves durante a decodificação. Ele opera em um sumário espectral local de página que é livre de treinamento (training-free) e integrado como um plugin plug-and-play para o vLLM.

Mecanismo Central

1. Decomposição Espectral Local de Página:

   • Em vez de usar uma base global ou de nível de sequência compartilhada, o LOCKS constrói uma base de baixo posto única para cada página do cache KV.
   • Quando uma página é finalizada (durante o prefill ou uma atualização de travessia de página), ela passa por uma pequena decomposição de autovalores de seus desvios de chave (chaves centralizadas).
   • O sumário armazena os $r$ autovetores superiores (base) e os coeficientes correspondentes para cada token da página.
   • Quantização: A base é armazenada em int4 e os coeficientes/centroides em int8. Com rank $r=8$ e tamanho de página $B=16$, o sumário é aproximadamente 10% dos bytes originais de KV da página.

2. Reconstrução e Seleção:

   • A cada etapa de decodificação, o LOCKS reconstrói os logits intra-página para a consulta usando o sumário compacto:
     $$\hat{s}_j(q) = \log \sum_{i \in P_j} \exp\left( \frac{q^\top \mu_j + (V_j^\top q)^\top c_i}{\sqrt{d}} \right)$$
     onde $\mu_j$ é o centroide da página, $V_j$ é a base local da página e $c_i$ são os coeficientes.
   • Ele calcula o Log-Sum-Exp (LSE) desses logits reconstruídos para estimar a massa de atenção da página.
   • A seleção é realizada classificando as páginas com base nessas massas estimadas. Crucialmente, nenhuma chave ou valor candidato é lido durante esta fase de seleção.

3. Seleção Compartilhada por Grupo (GQA):

   • Para a Atenção de Consulta Agrupada (GQA), o LOCKS faz a média das massas estimadas entre as cabeças de consulta em um grupo para formar uma única tabela de seleção compartilhada. Isso garante que todas as cabeças de um grupo atendam ao mesmo conjunto de páginas, evitando fragmentação de tráfego.

4. Implementação:

   • Implementado como um plugin para o vLLM inalterado.
   • Toda a etapa de decodificação (pontuação, seleção e atenção esparsa) roda dentro de grafos CUDA completos, garantindo overhead zero de despacho Python e permitindo a fusão eficiente de kernels.
   • A construção do sumário é desacoplada do caminho crítico de decodificação, ocorrendo apenas quando as páginas são finalizadas.

3. Principais Contribuições

Descobertas Teóricas

• Localidade da Estrutura: O artigo demonstra que a estrutura geométrica relevante para a seleção de páginas está concentrada no escopo da página, não no escopo da sequência ou do cache. A própria base de autovalores de rank-8 de uma página captura a maior parte da energia de sua chave, enquanto bases compartilhadas (mesmo quando combinadas em tamanho de estado) falham em capturar direções específicas da página.
• Impossibilidade de Projeções Compartilhadas: Os autores provam (Proposição 1) que qualquer projeção linear de baixo posto fixa e independente da consulta é cega ao conteúdo para uma família $(d-\rho)$-dimensional de diferenças intra-página. Isso significa que bases compartilhadas não conseguem distinguir entre páginas que diferem apenas em direções de conteúdo específicas, levando a erros de seleção.
• Garantias de Retenção: O artigo fornece um limite composto (Teorema 1) mostrando que o erro na saída da atenção é controlado diretamente pela energia da cauda espectral da reconstrução local da página. A base local da página é provada como residualmente ótima entre reconstruções ortogonais de rank-$r$.

Resultados Empíricos

• Recuperação e QA (LongBench-v1, RULER): O LOCKS acompanha o oráculo "ler-cada-chave" (seleção LSE exata) até os menores orçamentos. Em tarefas de recuperação densa (RULER), ele supera significativamente os seletores de última geração (Quest, KVzip, ShadowKV, RocketKV) em retenção de portadores (manter as páginas que contêm a resposta), que é o principal motor do sucesso da tarefa.
• Raciocínio de Longo Formato (AIME26, MATH-500): Em tarefas de raciocínio onde o modelo gera milhares de tokens, os seletores de base costumam colapsar em orçamentos apertados. O LOCKS mantém alta precisão, igualando o desempenho do FullKV em contextos de 100K+ enquanto atende a apenas ~2% dos tokens.
• Eficiência:
  • Com 1M de comprimento de contexto, o LOCKS alcança um ganho de velocidade de 2.0× na latência por token em comparação com a atenção densa (FlashAttention-3/FlashInfer).
  • Ele reduz os bytes de KV lidos por etapa de decodificação em 9.8× em 1M de contexto.
  • O ganho de throughput escala com o tamanho do lote, atingindo 1.80× em 256K de contexto com lotes grandes.

4. Significância e Alegações

O artigo afirma que a concentração espectral local de página é o elo perdido para a decodificação esparsa eficiente. Ao deslocar o escopo da representação do global/sequência para o local (por página), o LOCKS preserva as páginas "portadoras" necessárias para a conclusão da tarefa sem exigir treinamento ou modificar a arquitetura do modelo.

• Integração Plug-and-Play: Diferente de muitos trabalhos anteriores que exigem forks de modelos ou portões de treinamento, o LOCKS funciona com o vLLM inalterado e modelos GQA padrão.
• Alinhamento Teórico-Prático: O método é um dos poucos que conecta um limite teórico de erro de reconstrução diretamente à qualidade da tarefa final e à sobrevivência do portador, validando que a "corrente" da concentração espectral até o desempenho da tarefa se sustenta.
• Escalabilidade: A abordagem escala para contextos extremos (1M+ de tokens) onde a largura de banda da memória é o principal gargalo, oferecendo um caminho prático para servir modelos de contexto longo sem recorrer ao offloading para CPU ou memória em camadas.

Os autores mantêm modéstia quanto às limitações: a retenção do sumário quantizado implantado é medida empiricamente em vez de estritamente limitada por etapa, e o método atualmente assume um tamanho de página fixo. No entanto, os resultados estabelecem que um sumário espectral local de página, livre de treinamento, é suficiente para rastrear a seleção do oráculo até orçamentos de esparsidade agressivos, permitindo a aceleração prática de contexto longo.