Resumo Técnico: RKSC – Compartilhamento de Cache KV Consciente de Raciocínio e Saída Antecipada Confiante

1. Declaração do Problema

Modelos de linguagem de grande escala (LLMs) modernos projetados para raciocínio complexo (ex: modelos estilo DeepSeek-R1 ou o1) frequentemente empregam estratégias de inferência de múltiplos ramos, como Tree of Thoughts ou autoconsistência, onde múltiplas trajetórias de raciocínio são geradas para um único problema. O artigo identifica duas redundâncias estruturais em sistemas de inferência de produção atuais (como vLLM e SGLang) que degradam a eficiência nesses pipelines:

1. Desperdício de KV entre Ramos (Cross-Branch KV Waste): Ramos paralelos tipicamente compartilham uma parte significativa de seu prefixo (a declaração do problema e o contexto inicial). No entanto, os sistemas existentes apenas reutilizam o cache de Chave-Valor (KV) quando as sequências de tokens são byte a byte idênticas. Se os ramos divergirem semanticamente, mas compartilharem os mesmos tokens de prefixo, ou se compartilharem um prefixo mas utilizarem frases diferentes para dicas (hints), o sistema recomputa o custoso processamento de atenção $O(n^2)$ para o prefixo compartilhado de forma independente para cada ramo.
2. Verificação Excessivamente Profunda (Over-Deep Verification): Em pipelines de verificação de recompensa de processo, o modelo realiza uma passagem forward de profundidade total para pontuar cada ramo gerado, mesmo quando a confiança da geração já é decisiva ou quando o estado interno do modelo (distribuição de logits) se estabilizou em uma camada intermediária. Os sistemas atuais carecem de mecanismos para pular a passagem de verificação inteiramente ou para terminá-la precocemente com base na entropia por camada.

2. Metodologia: O Framework RKSC

Os autores propõem o RKSC (Reasoning-Aware KV Cache Sharing), um framework de inferência livre de treinamento que elimina essas redundâncias através de três mecanismos complementares e alternáveis. O RKSC não requer ajuste fino (fine-tuning), mudanças de arquitetura e nem kernels CUDA customizados.

A. ASKS (Attention-Similarity KV Sharing)

O ASKS generaliza o cache de prefixo por correspondência exata de tokens para um regime semântico.

• Mecanismo: Em vez de exigir tokens byte a byte idênticos, o ASKS computa uma similaridade de cosseno ponderada entre os estados ocultos do sufixo de um ramo e os estados ocultos da raiz.
• Ponderação: A pontuação de similaridade ($\sigma_b$) aplica pesos exponencialmente crescentes às camadas posteriores do transformer ($w_l = \exp(\alpha l/L)$), pois estas camadas carregam mais informações específicas da tarefa e são melhores em detectar divergência semântica.
• Ação: Se $\sigma_b \geq \tau$, o ramo é considerado semanticamente consistente com o prefixo e recebe o cache KV compartilhado via uma operação repeat_interleave de custo zero. Ramos abaixo do limiar retornam ao cálculo independente.
• Generalização: Isso generaliza estritamente o cache de tokens exatos (onde $\sigma_b=1$ para tokens idênticos) e permite a reutilização para ramos com diferentes formulações superficiais, mas representações semânticas similares.

B. CGEE (Confidence-Gated Early Exit)

O CGEE reduz os custos de verificação através de uma estratégia de saída de dois níveis:

1. Pular Verificação (Nível 1): Antes de executar a passagem de verificação, o sistema verifica a confiança da geração. Se o ramo vencedor possuir alta confiança absoluta ($p(b) \geq \tau_{conf}$) e dominar os competidores por uma margem relativa clara, toda a passagem forward de verificação é ignorada.
2. Saída por Entropia de Camada (Nível 2): Se a verificação prosseguir, hooks forward leves monitoram a entropia de logit em cada camada. Se a entropia cair abaixo de um limiar ($\theta$) e se estabilizar (mudança $< \epsilon$) em uma camada intermediária $l^*$, a passagem forward é terminada antecipadamente. O output parcial é usado como o resultado da verificação, pulando as camadas restantes $L - l^*$.

C. RSBCM (Reasoning-Selective Block Cache Manager)

Para evitar o crescimento descontrolado de memória durante buscas em árvores profundas, o RSBCM gerencia o cache compartilhado.

• Política de Expulsão: Blocos são expulsados com base em uma pontuação de importância $\omega = \text{branch\_score} / (\text{depth} + 1)$. Isso prioriza a manutenção de blocos de ramos com pontuações altas em profundidades rasas, enquanto expulsa primeiro os blocos de trajetórias profundas e de baixa confiança.

3. Principais Contribuições

• ASKS: Introduz um portão de similaridade semântica que alcança 28,6% de reutilização adicional de KV sobre a correspondência exata de tokens em cenários de frases diversas, generalizando estritamente o cache de prefixo existente.
• CGEE: Propõe um mecanismo de saída de dois níveis que pula a verificação quando a confiança é decisiva e trunca a passagem quando a entropia por camada se estabiliza, reduzindo a latência de verificação sem retreinamento.
• RSBCM: Fornece uma estratégia de expulsão baseada em prioridade de profundidade para manter a estabilidade do cache em árvores de busca profundas.
• Design de Custo Zero: O framework não adiciona parâmetros aprendíveis e não requer modificações arquiteturais, atuando como um wrapper de aplicação direta para pipelines de raciocínio existentes.

4. Resultados Experimentais

O framework foi avaliado em cinco famílias de modelos (7B–10B parâmetros, incluindo Qwen2.5, Mistral, Falcon3 e Llama-3) e quatro benchmarks (GPQA Diamond, MMLU-STEM, ARC-Challenge, GSM8K) em 1.000 problemas.

• Aceleração (Speedup): O RKSC alcança um aumento médio de velocidade de 3,008× sobre a linha de base No-KV (que recomputa prefixos) e uma melhoria de 1,66× sobre o cache de prefixo de tokens exatos equivalente ao vLLM. O pico de aceleração observado foi de 3,990× (Llama-3-8B no MMLU-STEM).
• Acurácia: O framework mantém alta acurácia com uma taxa de erro induzida pelo CGEE de apenas 0,37% (6 erros em 1.616 chamadas de verificação); os erros concentraram-se em benchmarks de múltipla escolha onde a fluência da geração ocasionalmente divergia de sinais de correção; tarefas de geração de texto livre mostraram taxa de erro próxima de zero.
• Decomposição por Componente:
  • O compartilhamento de KV sozinho proporcionou um aumento de velocidade de 1,343× (redução de latência de 25,5%).
  • O sistema combinado RKSC (KV + CGEE) proporcionou 1,622× de aceleração na decomposição de throughput e 3,008× em avaliação estendida.
  • A taxa de "pular verificação" do CGEE foi adaptável à dificuldade da tarefa, variando de ~42% em problemas difíceis do GPQA até ~92% em problemas mais fáceis do MMLU-STEM.
• Agnosticismo de Arquitetura: O ganho de velocidade no compartilhamento de KV foi consistente em todos os cinco modelos (desvio padrão de 0,02×), confirmando que os ganhos são impulsionados pela estrutura da carga de trabalho de inferência (comprimento do prefixo e fator de ramificação) e não por arquiteturas de modelos específicas.

5. Significância e Alegações

O artigo afirma que o RKSC aborda ineficiências estruturais fundamentais no raciocínio de LLMs de múltiplos passos que os sistemas de produção atuais ignoram. Ao desacoplar a reutilização de KV da identidade lexical (ASKS) e explorar o fenômeno de "colapso de entropia" em camadas posteriores do transformer (CGEE), o RKSC reduz significativamente os custos de energia e tempo da inferência escalonada.

Os autores enfatem que esses ganhos são alcançados sem o overhead computacional de treinamento ou a complexidade arquitetural de métodos de decodificação especulativa. O framework é apresentado como uma otimização prática e de aplicação direta que torna o raciocínio de última geração mais acessível em hardware restrito, reduzindo a computação redundante em pipelines de múltiplos ramos. Os resultados sugerem que, para tarefas de raciocínio de múltiplos passos, uma parte significativa da computação de verificação e de preenchimento (prefill) é estruturalmente redundante e pode ser com segurança eliminada ou encurtada com base em métricas de similaridade semântica e confiança.