Resumo Técnico: SIFT – Índice Seletivo para Computação Rápida de Prefill de RAG ao Explorar a Invariância da Atenção

1. Declaração do Problema

A Geração Aumentada por Recuperação (RAG) melhora as respostas de Grandes Modelos de Linguagem (LLMs) injetando documentos relevantes na consulta do usuário. No entanto, isso aumenta significativamente o comprimento do prompt ($L$), levando a um aumento quadrático ($O(L^2)$) no custo computacional do mecanismo de atenção durante a fase de prefill. Isso resulta em um alto Tempo para o Primeiro Token (TTFT), degradando a experiência do usuário e o rendimento do sistema.

Soluções existentes tentam mitigar isso pré-computando tensores Chave/Valor (KV) para documentos de RAG offline e reutilizando-os durante a inferência online. No entanto, essas abordagens enfrentam duas limitações críticas:

1. Degradação da Acurácia: A reutilização total de KV ignora a atenção cruzada (interações entre documentos), causando uma queda severa na acurácia. Trabalhos anteriores que tentam corrigir isso via recomputação seletiva de grão grosso (por exemplo, recomputando um 20% fixo de tokens em todas as camadas) falham em capturar os padrões de atenção diversos e de grão fino únicos para cada combinação de (documento, cabeça, camada). Por exemplo, o CacheBlend degrada a acurácia do Llama-8B no LongBench em até 68%.
2. Gargalos de Latência: Em GPUs modernas, a latência de transferência de grandes tensores KV do disco (SSD) para a memória frequentemente excede o tempo necessário para recomputá-los totalmente. À medida que as capacidades de computação das GPUs superam a largura de banda de armazenamento, os métodos de reutilização de KV tornam-se limitados por E/S (I/O), anulando potenciais ganhos de velocidade.

2. Metodologia: SIFT

O artigo propõe o SIFT (Índice Seletivo para Computação Rápida de Prefill de RAG ao Explorar a Invariância da Atenção). O SIFT evita armazenar dados KV inteiramente. Em vez disso, ele armazena apenas as localizações de altos escores de atenção como metadados compactos, permitindo uma recomputação seletiva de grão fino durante a fase de prefill online.

Insights Principais: Invariância da Atenção

O SIFT baseia-se em duas observações principais relativas à invariância dos padrões de atenção em documentos de RAG:

1. Invariância da Atenção Local: As localizações espaciais de altos escores de atenção dentro da autoatenção de um documento (atenção local) permanecem invariantes independentemente do contexto circundante (ou seja, quais outros documentos são prefixados). Embora os valores dos escores mudem com o contexto, as localizações dos altos escores não mudam.
2. Consistência da Atenção Cruzada: As chaves dentro de um documento que atraem altos escores de atenção durante o processamento isolado também tendem a atrair uma forte atenção cruzada de futuros documentos coexistentes. Tokens semanticamente importantes atraem atenção consistentemente.

Design do Sistema

• Processamento Offline: O SIFT realiza um prefill denso em cada documento de RAG independentemente para analisar as distribuições de escore de atenção. Ele identifica as localizações de alto escore e as codifica em dois vetores de bits compactos por documento (por camada e cabeça):
  • Vetor de Bits de Atenção Local (LA): Codifica tiles de alto escore dentro da matriz de autoatenção do documento.
  • Vetor de Bits de Atenção Cruzada (CA): Codifica páginas de KV (colunas) que consistentemente atraem alta atenção cruzada de outros documentos.
• Eficiência de Armazenamento: Em vez de armazenar MBs ou GBs de tensores KV, o SIFT armazena apenas alguns KB de vetores de bits. Isso reduz os requisitos de armazenamento em até 24.000× em comparação com os caches de KV, permitindo que os metadados residam inteiramente na rápida DRAM da CPU, eliminando gargalos de transferência de disco.
• Execução Online: Durante a fase de recuperação de RAG, os metadados do SIFT são buscados junto com os documentos top-k. Durante a fase de prefill, um kernel de atenção esparsa customizado lê esses vetores de bits. Ele computa os escores de atenção apenas para as localizações marcadas (tiles/páginas) e ignora o restante.

Detalhes de Implementação

• Kernel de Decodificação: Um kernel customizado decodifica os vetores de bits em índices de tiles explícitos com overhead negligenciável (~73 µs).
• Kernel de Atenção Esparsa: Estende o FlashAttention-3 para iterar apenas sobre os índices esparsos definidos pelos vetores de bits, realizando a recomputação seletiva.

3. Principais Contribuições

1. Identificação da Invariância da Atenção: O artigo estabelece que os padrões de esparsidade da atenção local são invariantes ao contexto circundante e que chaves de alta atenção intra-documento predizem alta atenção cruzada.
2. Arquitetura SIFT: Desenvolvimento de um sistema que codifica localizações de alta atenção em dois vetores de bits compactos, permitindo a recomputação seletiva de grão fino sem armazenar dados KV.
3. Desempenho e Acurácia: O SIFT alcança até 1,71× de aceleração de TTFT sobre a recomputação total, mantendo a acurácia média dentro de 1% da recomputação total. Ele supera significativamente as linhas de base de reutilização de KV (como o CacheBlend), que sofrem tanto com a degradação da acurácia quanto com a latência limitada pelo disco.

4. Resultados

Avaliado nos datasets LongBench (2WikiMQA, HotpotQA, TriviaQA, Musique) utilizando Llama-3.1-8B, MiniMax-M2.5 e Qwen3-235B-A22B em GPUs NVIDIA H200:

• Aceleração de TTFT: O SIFT supera consistentemente a recomputação total, com acelerações variando de 1,26× a 1,71×, dependendo do comprimento do contexto e do modelo. As acelerações aumentam com o comprimento do contexto, conforme a atenção se torna mais esparsa.
• Acurácia: O SIFT mantém a acurácia dentro de 1% da recomputação total em todas as configurações testadas. Em contraste, o CacheBlend mostra uma degradação de acurácia de até 68,2% no Llama-8B.
• Armazenamento e Latência: Os metadados do SIFT são ~24.000× menores que os caches de KV (ex: 0,98 MB vs. 15,1 GB para um contexto de 64K). Isso permite que os metadados sejam armazenados na DRAM, tornando a latência de transferência de disco negligenciável (<0,11% do TTFT).
• Eficiência Energética: O SIFT alcança o melhor Produto Energia-Atraso (EDP) e o melhor Produto Energia-Atraso ao Quadrado (ED²P), sendo 3,1× melhor que a recomputação total e 3,2× melhor que o CacheBlend, devido à redução de computação e eliminação de E/S de disco.

5. Significância

O artigo afirma que o SIFT aborda o compromisso fundamental entre redução de computação e acurácia no prefill de RAG. Ao explorar a propriedade de reutilização de contexto das cargas de trabalho de RAG e a invariância de atenção dos documentos, o SIFT possibilita uma estratégia de recomputação de grão fino que é simultaneamente eficiente computacionalmente e leve em termos de armazenamento.

Diferente de trabalhos anteriores que dependem de heurísticas de grão grosso ou sofrem com a "parede de disco" da transferência de massivos caches de KV, o SIFT fornece uma solução escalável para arquiteturas de GPU modernas onde a computação é rápida, mas a largura de banda de memória é o gargalo. Os autores afirmam que o SIFT permite que sistemas de RAG alcancem uma acurácia próxima à da recomputação total com latência e consumo de energia significativamente reduzidos, tornando-o uma abordagem viável para sistemas de produção de RAG que exigem baixo TTFT.