Resumo Técnico: Cache de Prefixo Esparso para Atendimento de LLMs Híbridos e Recorrentes

1. Declaração do Problema

Os sistemas existentes de atendimento de Modelos de Linguagem de Grande Escala (LLMs) dependem fortemente de cache de prefixo denso, onde os estados Chave/Valor (KV) para cada token em um prefixo compartilhado são armazenados para evitar computação de preenchimento redundante. Esta abordagem é padrão para arquiteturas Transformer. No entanto, Modelos de Espaço de Estados (SSMs) e arquiteturas híbridas (combinando atenção com camadas recorrentes) apresentam uma assimetria estrutural: uma camada recorrente evolui seu estado como $h_t = F(h_{t-1}, x_t)$, o que significa que pode retomar a computação a partir de um único estado armazenado $h_c$, em vez de exigir todo o histórico de estados intermediários.

O cache denso por token de estados recorrentes é frequentemente desnecessário e proibitivamente caro. Para modelos como o Qwen-3.5, o estado recorrente por cabeça escala como $d_{head}^2$, enquanto o armazenamento KV de atenção escala linearmente. Consequentemente, armazenar o estado recorrente de cada token consome significativamente mais memória do que os caches KV de atenção, sem oferecer benefício adicional sobre o armazenamento de um conjunto esparso de pontos de verificação.

O problema central abordado é: Dado um orçamento de memória de $M$ slots de pontos de verificação para um prefixo de comprimento $N$, onde os pontos de verificação devem ser posicionados para minimizar o custo de recomputação esperado sob uma distribuição específica de profundidades de sobreposição de solicitações futuras? Esta otimização é distinta de políticas de admissão/evicção de cache; foca em quais tokens verificar dentro de uma única entrada retida para maximizar a utilidade para sobreposições parciais de prefixo.

2. Metodologia

Formulação Teórica

Os autores formalizam o problema como um problema de mediana k ponderada unilateral em uma linha.

• Seja $C = \{c_1, \dots, c_M\}$ o conjunto de posições de pontos de verificação.
• Para uma solicitação que sobrepõe o prefixo até a profundidade $t$, o ponto de verificação reutilizável mais profundo é $\ell(t; C) = \max(\{0\} \cup \{c \in C : c \le t\})$.
• O custo é o número de tokens a recomputar: $r(t; C) = t - \ell(t; C)$.
• O objetivo é minimizar o custo de recomputação esperado $E[r(T; C)]$ dada uma distribuição de profundidade de sobreposição $T$ com probabilidades $p_t$.

Abordagem Algorítmica

1. Posicionamento Balanceado Ótimo (Uniforme/Pior Caso): Os autores provam que, sob uma distribuição de sobreposição uniforme ou no pior cenário, o espaçamento balanceado (pontos de verificação distribuídos uniformemente) é ótimo.
2. Programação Dinâmica (DP) Consciente da Distribuição: Para distribuições de sobreposição não uniformes e do mundo real, os autores derivam um algoritmo exato de programação dinâmica $O(NM)$.
   • A relação de recorrência é: $dp[m, j] = \min_{1 \le s \le j} (dp[m-1, s-1] + w(s, j))$, onde $w(s, j)$ é o custo de usar $s$ como o ponto de verificação mais profundo para profundidades $s \dots j$.
   • Isso é resolvido eficientemente usando o truque do casco convexo monótono, aproveitando a monotonicidade das somas de prefixo.
3. Tratamento de Deriva de Distribuição: Como a verdadeira distribuição de sobreposição é desconhecida e deve ser estimada a partir de dados históricos, o artigo fornece garantias teóricas (via Lema 1 e Teorema 3) de que a função valor é Lipschitz contínua em relação à distribuição de sobreposição. Isso garante que o uso de um histograma empírico (ou um histograma ponderado exponencialmente para distribuições em deriva) produz um agendamento quase ótimo com subotimalidade limitada.

Estratégia de Implementação

• Checkpointing: O sistema armazena estados recorrentes exatos em posições esparsas. Ao ocorrer um acerto de cache, o sistema retoma a partir do ponto de verificação armazenado mais profundo e recompute o sufixo exatamente.
• Exatidão: Como a atualização recorrente é determinística, retomar de um ponto de verificação e replayar o sufixo produz saídas idênticas em bits a um preenchimento completo.
• Integração Híbrida: Para modelos híbridos, o método armazena o cache KV de atenção completo (denso), mas aplica checkpointing esparso apenas às camadas recorrentes.

3. Contribuições Principais

1. Optimalidade Teórica: Provou-se que o espaçamento balanceado é ótimo sob sobreposição uniforme e condições de pior caso (Teorema 1).
2. Algoritmo Exato: Formulou-se o cache de prefixo esparso como um problema de mediana k ponderada e forneceu-se uma solução exata de programação dinâmica $O(NM)$ (Teorema 2).
3. Garantias de Robustez: Demonstrou-se que o agendamento ótimo é estável sob especificação incorreta da distribuição. A subotimalidade é limitada pela distância de variação total entre as distribuições de sobreposição empírica e verdadeira (Teorema 3), e estendida a cenários com deriva de distribuição usando reponderação exponencial (Teorema 4).
4. Validação Empírica: Validou-se o método em conjuntos de dados do mundo real (QuALITY, NarrativeQA, System Prompts), mostrando melhorias consistentes sobre baselines de orçamento fixo.

4. Resultados Experimentais

Conjuntos de Dados e Configuração

• Conjuntos de Dados:
  • QuALITY: Documento único com múltiplas perguntas (alta sobreposição perto do comprimento total).
  • NarrativeQA: Estrutura similar com ampla sobreposição multimodal.
  • System Prompts: Prompts de sistema reais combinados com consultas de usuário (sobreposição concentrada em comprimentos curtos).
• Baselines: Comparado contra "Sem Cache", "Apenas KV", "Balanceado" (espaçamento uniforme), "Bloco" (tamanho de bloco fixo), "Logarítmico" e estratégias com espaçamento "$\sqrt{L}$".
• Hardware: Execuções prototípicas em NVIDIA RTX 2080 Super com o modelo Qwen-3.5-0.8B.

Descobertas

• Economia de Tokens: O agendamento ótimo via DP domina consistentemente as baselines determinísticas (Balanceado, Bloco, Logarítmico) em termos de fatores de redução de trabalho recorrente.
• Vantagem em Orçamento Baixo: A lacuna de desempenho é maior em orçamentos baixos de pontos de verificação, onde a distribuição de sobreposição é mais não uniforme. Por exemplo, em System Prompts (onde as sobreposições são curtas), o método DP posiciona pontos de verificação densamente perto do pico de divergência, alcançando o mesmo tempo de execução que o cache de bloco com menos pontos de verificação.
• Tempo de Parede (Wall-Clock Time): Em QuALITY e System Prompts, o agendamento derivado via DP produz acelerações mensuráveis de tempo de parede em grupos de camadas representativos. A redução no tempo de execução correlaciona-se linearmente com o número de tokens economizados.
• Sobrecarga: A sobrecarga de carregar pontos de verificação da CPU para a GPU é pequena em relação às economias obtidas pela redução da recomputação, resultando em uma melhoria líquida no tempo de execução.

5. Significado e Alegações

O artigo alega que o cache de prefixo esparso desbloqueia um novo ponto de design para o atendimento de LLMs híbridos e recorrentes que anteriormente não estava disponível.

• Compromisso Memória/Computação: Oferece um trade-off entre uso de memória e recomputação que o cache KV denso não pode fornecer, especificamente para arquiteturas onde o estado pode ser extraído e restaurado exatamente.
• Exatidão: Diferentemente de técnicas de compressão KV aproximadas, este método preserva saídas exatas sem alterar os kernels de computação recorrente.
• Aplicabilidade: É mais benéfico para cargas de trabalho onde as solicitações compartilham um prefixo substancial, mas não idêntico (por exemplo, RAG com perguntas diferentes, prompts de sistema variados). Não oferece benefícios adicionais para cargas de trabalho de chat estritamente de apenas anexação (onde o último estado é suficiente), mas não incorre em sobrecarga nesses casos.
• Trabalho Futuro: Os autores reconhecem desafios de engenharia, especificamente a necessidade de APIs eficientes de extração/restauração de estado em tempos de execução de produção e a integração desta estratégia de posicionamento com políticas de admissão/evicção (por exemplo, Marconi et al.). Eles também observam que a teoria atual de prefixo único poderia ser estendida para árvores de prefixo gerais.