Resumo Técnico: ParisKV

Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) estão estendendo rapidamente suas janelas de contexto para milhões de tokens, criando uma pressão sem precedentes na eficiência da inferência. O cache KV, que armazena chaves e valores para a decodificação autorregressiva, domina rapidamente a pegada de memória e a latência conforme o contexto cresce. A decodificação de contexto longo é inerentemente limitada pela memória, pois cada etapa requer a leitura de um grande volume de vetores KV, com custos de largura de banda que escalam linearmente com o comprimento do contexto.

Métodos de atenção esparsa/seletiva tentam mitigar isso ao atender apenas a um subconjunto Top-k de tokens passados relevantes, mas as abordagens existentes de recuperação de cache KV enfrentam três desafios críticos em cenários de geração longa e sensíveis à latência:

1. Compromisso Velocidade–Qualidade: A recuperação aproximada leve (ex: agrupamento grosseiro, quantização de baixa precisão) frequentemente sacrifica o recall em favor da velocidade. Recuperar a precisão geralmente exige aumentar o orçamento de recuperação, o que erode os benefícios da esparsidade.
2. Deriva de Decodificação (Decoding Drift): Centróides aprendidos a partir do agrupamento de chaves de preenchimento (prefill) tornam-se obsoletos conforme a geração continua. À medida que as chaves de decodificação se acumulam, a distribuição muda, causando um descompasso entre os centróides do prefill e a verdadeira distribuição de chaves. Isso leva a um colapso severo de recall durante etapas de decodificação longas em métodos anteriores.
3. Recuperação Limitada pela CPU: Quando os caches KV completos são transferidos para a memória da CPU para economizar espaço na GPU, a recuperação frequentemente envolve busca no lado da CPU e transferência de dados CPU-para-GPU. Isso introduz erros de aproximação e torna a orquestração da CPU o fator dominante na latência de ponta a ponta.

Metodologia: O Framework ParisKV

ParisKV é um algoritmo–sistema de co-design consciente de hardware para recuperação de cache KV que alcança decodificação rápida e recall estável sob entrada longa e/ou geração longa (escala de até um milhão de tokens), mesmo na presença de deriva de distribuição.

Inovações Algorítmicas Principais

1. Centróides Analíticos Robustos à Deriva:
Ao contrário de métodos anteriores que aprendem centróides a partir de chaves de prefill, o ParisKV transforma consultas e chaves em um espaço estável e independente de dados.

• Normalização e Rotação: Chaves e consultas são $\ell_2$-normalizadas e submetidas a uma rotação ortogonal aleatória compartilhada (implementada via SRHT). Isso as mapeia em uma hiperesfera unitária.
• Centróides Analíticos: Em vez de aprender centróides, o ParisKV define um conjunto de centróides uniformemente distribuídos na hiperesfera unitária baseados em padrões de sinal (ex: $\pm 1/\sqrt{m}$). Como esses centróides são independentes de dados e uniformemente distribuídos, qualquer nova chave gerada durante a decodificação estará sempre próxima de pelo menos um centróide, prevenindo o problema de "centróide obsoleto".

2. Pipeline de Recuperação Nativo de GPU em Duas Etapas:
O ParisKV implementa um pipeline de recuperação de fase grossa-para-fina inteiramente na GPU para evitar gargalos de CPU.

• Geração de Candidatos Grosseiros (Baseada em Colisão): As chaves rotacionadas são divididas em $B$ subespaços. Cada direção de subespaço é atribuída ao centróide analítico mais próximo. Para uma consulta, o sistema seleciona os top-$\rho$ centróides mais similares em cada subespaço. As chaves nesses buckets selecionados recebem votos (com pesos em níveis). O sistema soma os votos através dos subespaços para gerar uma pontuação inteira grosseira, podando o espaço de busca para um pequeno conjunto de candidatos (ex: top 5–10%).
• Reclassificação de Candidatos (RSQ-IP): Os candidatos são reclassificados usando um estimador de produto interno calibrado sem acessar chaves de precisão total.
  • Quantização: Cada direção de subespaço é codificada usando um código de 4 bits (1 bit de sinal + 3 bits de magnitude) derivado de priors Beta induzidos pela rotação.
  • Calibração: Um fator de escala leve ($w_{i,b}$) é pré-computado durante a fase de prefill para corrigir erros de alinhamento de quantização.
  • Estimativa: O sistema estima o score de atenção bruto $\langle k, q \rangle$ usando a fórmula: $\|q\|_2 \sum_{b=1}^B w_{i,b} \langle v_{i,b}, \tilde{q}_b \rangle$, onde $v_{i,b}$ é a direção desquantizada e $\tilde{q}_b$ é a consulta rotacionada. Isso permite uma reclassificação de alta fidelidade usando apenas metadados compactos.

Inovações de Sistema

1. Offloading via Endereçamento Virtual Unificado (UVA):
O ParisKV transfere o cache KV de precisão total para a memória da CPU, mas utiliza UVA para permitir que os kernels da GPU busquem diretamente apenas os pares KV Top-k selecionados sob demanda. Isso evita operações explícitas de memcpy e overhead de agendamento do lado da CPU, permitindo que a GPU acesse a memória da CPU como se fosse local.

2. Mecanismo de Atualização de Janela Deslizante:
Para suportar geração ilimitada, o cache KV é gerenciado em quatro regiões: Sink (tokens iniciais na GPU), Local (tokens recentes na GPU), Retrieval (tokens históricos transferidos para CPU) e um Buffer de Atualização. Quando o buffer enche, os tokens locais mais antigos são ejetados para a zona de Retrieval (transferidos para CPU), e os tokens do buffer são promovidos para Local. Isso garante que os metadados de recuperação permaneçam atualizados enquanto preserva o acesso rápido tanto ao contexto recente quanto ao histórico de longo alcance.

Principais Contribuições

1. Centróides Analíticos Robustos à Deriva: Uma transformação independente de dados (normalização + rotação) que mantém o recall de recuperação estável mesmo sob mudanças significativas de distribuição na escala de um milhão de tokens, eliminando a necessidade de centróides aprendidos que degradam com o tempo.
2. Pipeline de Fase Grossa-para-Fina Nativo de GPU: Um sistema de recuperação combinando poda baseada em colisão e reclassificação calibrada (RSQ-IP) que alcança seleção Top-k de alta fidelidade diretamente em códigos comprimidos sem busca no lado da CPU.
3. Offloading Escalável Baseado em UVA: Uma implementação que permite a busca de KV sob demanda com intervenção mínima da CPU, permitindo que o sistema escale para contextos de um milhão de tokens onde a atenção total esgota a memória.
4. Avaliação Abrangente: Demonstrou desempenho em benchmarks de raciocínio de geração longa (GPQA-Diamond, MATH500, AIME25) e compreensão de contexto longo (LongBench-V2, RULER).

Resultados Experimentais

Os autores avaliaram o ParisKV em benchmarks de raciocínio de geração longa (GPQA-Diamond, MATH500, AIME25) e compreensão de contexto longo (LongBench-V2, RULER) através de múltiplas famílias de modelos (Qwen-3-8B, DeepSeek-R1-Llama-8B, etc.).

Acurácia:

• O ParisKV iguala ou supera a qualidade da atenção total em benchmarks de entrada longa e geração longa.
• No AIME25, ele melhora a acurácia pass@8 em +40 a +77 pontos em relação aos baselines de estado da arte MagicPIG e PQCache.
• Em 7 de 9 configurações, o ParisKV iguala ou excede a acurácia da atenção total, superando significativamente os baselines de recuperação anteriores.

Eficiência:

• Throughput: O ParisKV alcança até 2.8× mais throughput que a atenção total dentro de seu intervalo de execução.
• Latência em Escala: Na escala de um milhão de tokens (1M tokens), o ParisKV reduz a latência de decodificação em 17× comparado ao MagicPIG e 44× comparado ao PQCache.
• Escalabilidade: Enquanto a atenção total sofre OOM (Out of Memory) com batch sizes $\ge 4$ para contextos de 128K e $\ge 2$ para 256K, o ParisKV escala para batch sizes de 8 e 5, respectivamente. Ele permanece executável em contextos de 384K e 1024K, onde a atenção total falha completamente.

Significância e Alegações

O artigo afirma que o ParisKV aborda os desafios fundamentais de deriva de distribuição e latência de recuperação na inferência de LLMs de contexto longo. Ao transformar as chaves em uma hiperesfera unitária e empregar centróides analíticos independentes de dados, o sistema elimina a obsolescência de centróides que afeta os métodos existentes.

Os autores afirmam que a recuperação de cache KV cuidadosamente projetada pode simultaneamente melhorar tanto a eficiência quanto a acurácia para a inferência de contexto longo, mesmo na presença de deriva. Isso abre novas possibilidades para o deploy de LLMs em aplicações de geração longa sensíveis à latência (como tarefas de raciocínio complexo) onde métodos anteriores ou falhavam em escalar ou sofriam de severa degradação de acurácia. O trabalho demonstra que é possível alcançar qualidade quase sem perdas mantendo os benefícios de velocidade da atenção esparsa, mesmo quando o cache KV é transferido para a memória da CPU.