Resumo Técnico: Spherical KV

1. Definição do Problema

A inferência de contexto longo em Grandes Modelos de Linguagem (LLMs) é cada vez mais limitada pelo cache de Chave-Valor (KV), e não pela vazão aritmética. À medida que o comprimento do contexto ($T$) cresce, a pegada de memória residente escala linearmente com $T$, tamanho de lote ($B$), camadas ($L$) e cabeças ($H$). Em ambientes de serviço de produção que utilizam memória paginada (ex: PagedAttention), o gargalo dominante desloca-se das Operações de Ponto Flutuante (FLOPs) para o tráfego de Memória de Alta Largura de Banda (HBM). Especificamente, a decodificação torna-se limitada pelo streaming repetido de páginas KV históricas da memória do dispositivo para as unidades de computação.

Soluções existentes tentam mitigar isso através de três alavancas:

1. Retenção/Evicção: Armazenar menos tokens (ex: StreamingLLM, H2O).
2. Quantização/Compressão: Armazenar menos bits (ex: KVQuant, KIVI).
3. Offloading: Armazenar o KV em outro lugar (ex: CPU/Armazenamento).

No entanto, esses métodos frequentemente falham em abordar plenamente o gargalo do caminho crítico. Muitas técnicas de compressão exigem a reconstrução de vetores densos a partir de representações comprimidas antes de computar os logits de atenção (um caminho de "descomprimir-então-produto-escalar"). Essa reconstrução reintroduz o tráfego de HBM e impede a fusão de kernels, anulando a economia teórica de memória em substratos de serviço paginados realistas. Além disso, políticas de evicção simples muitas vezes carecem de granularidade para preservar informações críticas (ex: evidências recuperadas) enquanto descartam estados de baixa utilidade.

2. Metodologia: Spherical KV

O artigo propõe o Spherical KV, um método de inferência de contexto longo que trata a alocação de KV como um problema de taxa-distorção fundamentado na geometria da atenção. Ele consiste em dois componentes ortogonais projetados para funcionar dentro de um substrato de serviço irregular e paginado: Angle-Domain Attention (ADA) e Rate–Distortion Retention (RDR).

2.1 Angle-Domain Attention (ADA)

O ADA substitui o primitivo padrão "descomprimir-então-produto-escalar" por uma abordagem nativa de computação que opera diretamente sobre códigos esféricos compactos.

• Parametrização Esférica: As chaves são decompostas em um raio escalar ($r = \|k\|$) e ângulos direcionais ($\phi = \Phi(\hat{k})$).
• Similaridade no Domínio Comprimido: Em vez de cachear vetores densos, o ADA armazena raios quantizados ($\tilde{r}$) e códigos angulares empacotados ($c_\theta$). Durante o loop quente de decodificação, o kernel de atenção computa os logits diretamente a partir desses códigos:
  $$\hat{\ell}(q, k) \approx \|q\| \cdot \tilde{r} \cdot f_b(c_\theta^q, c_\theta^k)$$
  onde $f_b$ é uma operação de kernel fundido que computa a similaridade de cosseno a partir dos códigos angulares.
• Invariante de Não Reconstrução: A restrição de design crítica é que o kernel de decodificação nunca materializa um vetor de chave denso ($\tilde{k} \in \mathbb{R}^d$) em HBM. Isso garante que as reduções nos bytes de KV armazenados se traduzam diretamente em redução de tráfego de HBM por token.

2.2 Rate–Distortion Retention (RDR)

O RDR é um mecanismo de controle que otimiza conjuntamente quais estados reter e em qual precisão (nível/tier) armazená-los, sujeito a um orçamento de memória estrito.

• Otimização Conjunta: Diferente dos métodos que tratam evicção e quantização separadamente, o RDR resolve um problema de alocação de taxa-distorção discreto. Para cada estado de token/cabeça/camada $i$, ele seleciona uma decisão de residência binária ($z_i \in \{0, 1\}$) e um nível de precisão ($b_i \in \mathcal{B}$).
• Surrogate de Distorção: O controlador utiliza um proxy de distorção mecanístico derivado da análise de erro do ADA. Ele estima que o erro de logit é dominado por erros de aproximação angular ($\epsilon_\theta$) escalonados pelas magnitudes da query e da chave.
• Restrições de Proteção: Para evitar falhas catastróficas em informações críticas (ex: evidências recuperadas, spans de segurança, outliers), o RDR impõe restrições de proteção rígidas ($I_{prot}(i)=1 \implies z_i=1, b_i=b_{max}$).
• Páginas Homogêneas por Nível: O controlador agrupa estados retidos em páginas onde todos os itens compartilham o mesmo nível de precisão. Isso permite leituras de streaming de passagem única, guiadas por cabeçalhos, evitando gathers irregulares.

2.3 Integração de Sistema

O Spherical KV opera em um substrato de KV paginado/irregular padrão (compatível com layouts do tipo vLLM). O pipeline mapeia chaves densas para códigos esféricos, aplica o RDR para filtrá-los e nivelá-los, e armazena o resultado em um formato comprimido. O kernel de decodificação consome esses códigos diretamente, mantendo uma invariante de "não-reconstrução-densa".

3. Principais Contribuições

1. Angle-Domain Attention (ADA): Um primitivo de nível de kernel que computa logits de atenção diretamente de códigos esféricos compactos, eliminando a penalidade de tráfego de HBM associada à reconstrução de chaves densas.
2. Rate–Distortion Retention (RDR): Um controlador unificado que aloca conjuntamente residência e níveis de precisão sob um orçamento fixo, priorizando estados de alta utilidade e protegendo os frágeis.
3. Avaliação Orientada à Implantação: O método é avaliado não em caches contíguos lúdicos, mas em um substrato de serviço paginado/irregular realista, medindo throughput de ponta a ponta, tráfego de HBM e estabilidade comportamental.
4. Deslocamento da Fronteira de Pareto: O artigo demonstra que o Spherical KV desloca a fronteira memória–qualidade–vazão, alcançando maior throughput com menores pegadas de memória em comparação ao KV denso e baselines fortes.

4. Resultados Experimentais

Os autores avaliaram o Spherical KV no Llama-3.1-8B-Instruct, Qwen2.5-14B-Instruct e gpt-oss-20b através de comprimentos de contexto de 8K, 32K e 128K. As cargas de trabalho incluíram modelagem de linguagem (PG-19), QA com foco em recuperação (LongBench: HotpotQA, 2WikiMultiHopQA) e rollouts de agentes (AgentBench, ToolBench).

• Eficiência de Throughput e Memória: Sob restrições de qualidade equivalentes (iso-quality), o Spherical KV melhorou o throughput de decodificação em 1,55×–1,72× enquanto reduziu os bytes de KV residentes por token em 24–42%.
• Escalabilidade de Contexto: Os ganhos foram mais pronunciados no regime de 128K, onde a residência de KV e o tráfego de HBM dominam o desempenho do sistema.
• Eficiência Realizada via Kernel: O método reduziu com sucesso os bytes de HBM por token gerado, confirmando que as economias de memória não foram anuladas por sobrecargas de reconstrução.
• Estabilidade: O sistema demonstrou degradação suave sob orçamentos mais apertados e preservou trajetórias de raciocínio e sequências de chamadas de ferramentas em cargas de trabalho de agentes, evitando as falhas abruptas observadas em baselines de evicção rígida.

5. Significância e Alegações

O artigo afirma que o Spherical KV aborda o gargalo fundamental da inferência de contexto longo: o conflito entre capacidade de memória e largura de banda de memória. Ao tratar a compressão de KV como um problema de computação nativa (onde a representação comprimida é diretamente executável pelo kernel), em vez de apenas um problema de armazenamento, o método alcança um deslocamento genuíno na fronteira operacional.

Os autores enfatizam que:

• A compressão deve ser computacionalmente nativa: Apenas armazenar menos bits é insuficiente se o caminho de decodificação exigir reconstrução densa.
• A memória deve ser orçada onde importa: O controle de taxa-distorção permite que o sistema gaste precisão em estados de alto impacto enquanto economiza em outros de baixa utilidade.
• A avaliação deve ser fiel à implantação: Os ganhos devem sobreviver a layouts paginados realistas, sobrecargas de metadados e restrições de kernel para serem considerados reais.

O Spherical KV é apresentado como um mecanismo para reduzir a residência de KV e o tráfego de HBM, preservando a eficiência de decodificação e a estabilidade comportamental, visando especificamente os regimes limitados por memória dos LLMs modernos.