Resumo Técnico: Prefill Esparso Orientado por Oráculo para Camadas Full/GQA em Modelos de Contexto Longo Híbridos

1. Declaração do Problema

O prefill de contexto longo permanece um gargalo computacional significativo no atendimento de modelos de linguagem de grande escala (LLMs). Embora as arquiteturas modernas adotem cada vez mais designs híbridos (combinando componentes locais, esparsos, lineares e recorrentes), as camadas de atenção total ou Grouped-Query Attention (GQA) frequentemente persistem como o principal mecanismo para recuperar e fundir evidências de longo alcance. Embora o GQA reduza a memória do cache KV ao compartilhar cabeças de chave-valor, ele não elimina a busca densa sobre os tokens históricos dentro da dimensão da sequência.

A questão diagnóstica central abordada neste relatório técnico é: Sob uma granularidade de suporte e um orçamento top-k fixos, quanto da busca de atenção densa nas camadas full/GQA existentes é realmente necessária para preservar o comportamento ao nível da tarefa?

Os métodos atuais de atenção esparsa (ex: Quest, DuoAttention, MInference, NSA) enfrentam um desafio de interpretação: uma queda na qualidade ao introduzir esparsidade pode decorrer de um orçamento insuficiente, uma seleção de indexador deficiente, compartilhamento de suporte grosseiro ou diferenças de implementação de tempo de execução. Sem uma referência controlada, torna-se difícil distinguir entre a viabilidade teórica do suporte esparso e as limitações práticas de um indexador ou sistema de tempo de execução.

2. Metodologia

2.1 O Oráculo Top-k de Massa de Atenção

Os autores introduzem um oráculo top-k de massa de atenção como um instrumento científico controlado, não como um método de aceleração implantável. Seu propósito é separar a "viabilidade do orçamento esparso" do "erro do indexador" e dos "efeitos de realização do tempo de execução".

• Mecanismo: Para cada camada e posição de query, o oráculo calcula a distribuição de atenção densa completa. Em seguida, agrega as pontuações de atenção através das cabeças de query para formar uma distribuição de importância de token média por cabeça ($P_{\ell,t}$).
• Seleção: Com base em um orçamento top-k especificado ($k$) e granularidade de suporte (por-query ou compartilhada por bloco de seleção), o oráculo seleciona os índices de tokens top-$k$ desta distribuição.
• Execução: O modelo recomputa a saída de atenção restrita apenas a este conjunto de suporte selecionado, utilizando as queries, chaves e valores originais. Crucialmente, o oráculo utiliza a atenção densa apenas para selecionar o suporte; a trajetória subsequente do modelo é computada de forma esparsa.
• Granularidade: O estudo avalia duas granularidades:
  • Por-query ($b_{sel}=1$): Cada posição de query seleciona seu próprio suporte.
  • Compartilhada por bloco de seleção ($b_{sel}=64$): Um único conjunto de suporte é selecionado para um bloco de 64 queries consecutivas para amortizar custos de sistema.

2.2 Indexador Colapsado por Cabeça via Destilação KL

Guiado pelo oráculo, os autores derivam um indexador de tempo de inferência para backbones GQA congelados.

• Princípio de Design: O modelo de backbone permanece inalterado. Um indexador auxiliar é treinado para prever a distribuição de importância de token ($P_{\ell,t}$) definida pelo oráculo.
• Objetivo de Treinamento: O indexador é treinado via destilação de divergência KL a partir da distribuição de massa de atenção densa.
• Arquitetura: Os autores propõem um design de colapso de cabeça. Em vez de manter canais de pontuação separados para cada cabeça de query ou grupo KV, o indexador colapsa o espaço de pontuação auxiliar entre as cabeças, prevendo uma distribuição de importância de token compartilhada. Isso reduz as ativações de projeção e elimina a necessidade de uma dimensão de cabeça auxiliar explícita.
• Inferência: O indexador pontua as chaves candidatas, seleciona os índices top-$k$, e a camada GQA esparsa recomputa a atenção apenas nestes tokens selecionados. O orçamento esparso ($k$) permanece ajustável durante a inferência sem necessidade de retreinamento.

2.3 Decomposição de Evidência

O artigo categoriza explicitamente os resultados em três tipos distintos de evidência para evitar a confusão entre viabilidade e prontidão de implantação:

1. Gap do Oráculo: A diferença entre a atenção densa e a atenção esparsa do oráculo (usando o suporte selecionado pelo oráculo). Isso testa se a família de suporte é teoricamente viável.
2. Gap do Indexador: A degradação adicional ao substituir o suporte do oráculo por um indexador destilado.
3. Gap de Realização: A degradação adicional introduzida por restrições de tempo de execução (ex: compartilhamento de seleção por bloco, fusão de kernel).

3. Principais Resultados

3.1 Viabilidade do Oráculo (Gap do Oráculo)

• Desempenho Próximo ao Denso: Em todo o conjunto de checkpoints da família Qwen (Qwen3-8B, Qwen3.5/3.6 27B/35B-A3B), o oráculo por-query permanece dentro de 1 ponto percentual da atenção densa nos contextos mais longos avaliados (até 128K).
• Sensibilidade ao Orçamento: Em um sweep da Qwen3.5-9B de 4K a 100K, o gap do oráculo manteve-se dentro de 0,48 pontos do denso sob um cronograma top-k configurado (256–2048), com esparsidade variando de 87,9% a 96,0%.
• Impacto da Granularidade: Embora os oráculos por-query sejam robustos, os oráculos compartilhados por bloco de seleção ($b_{sel}=64$) mostram um efeito de "abismo" (cliff effect). Em 128K, orçamentos altos (top-k=2048) mantêm qualidade próxima à densa, mas orçamentos baixos (top-k=256) causam colapso significativo em tarefas de recuperação intensiva. Isso indica que o compartilhamento de suporte requer escalonamento de orçamento consciente do comprimento.

3.2 Qualidade do Indexador Destilado (Gap do Indexador)

• Preservação de Qualidade: Usando indexadores destilados separadamente para Qwen3.5-0.8B e Qwen3.5-9B, os autores relatam nenhuma degradação agregada material em fatias de validação mantidas (contextos de 16K/32K) nos benchmarks RULER, VideoMME, MMLongBench-Doc e LongBench-v2.
• Gaps Macroscópicos: Os gaps macro reportados (Destilado menos Denso) foram de +2,04 para Qwen3.5-0.8B e +1,13 para Qwen3.5-9B. Os autores tratam esses gaps positivos como preservação de qualidade dentro da resolução de avaliação, não como melhorias intrínsecas sobre a atenção densa.
• Por-Query vs. Fundido: Em uma validação pura de RULER-32K, o caminho destilado por-query ficou dentro de 0,6 pontos do denso, enquanto o caminho fundido (compartilhamento por bloco de seleção) mostrou um gap de 5,0 pontos, destacando o custo da amortização amigável ao sistema.

3.3 Folga de Tempo de Execução (TTFT)

• Acelerações do Indexador Destilado: Medições em cartão único mostram acelerações de TTFT para o prefill esparso com indexador destilado:
  • Qwen3.5-0.8B (NPU): Até 1,71× de aceleração em contexto de 256K.
  • Qwen3.5-9B (GPU): Até 1,93× de aceleração em contexto de 256K.
  • As acelerações aumentam com o comprimento do contexto, conforme o custo relativo das camadas de não-atenção e overheads diminui.
• Testes de Estresse com Indexador Não Treinado: Executar o caminho esparso com pesos de indexador aleatoriamente inicializados (não treinados) revelou um potencial de folga ainda maior (até 3,44× no NPU em 128K), indicando que o próprio caminho de execução esparsa é eficiente, embora a qualidade da saída não seja validada nesta configuração.

4. Significância e Alegações

O artigo posiciona-se como um relatório técnico diagnóstico, não como uma reivindicação de um sistema esparso totalmente implantado e sem perdas. Suas principais contribuições são:

1. Reenquadramento do Problema: Ele enquadra o prefill esparso para camadas GQA como uma questão de redutibilidade orientada por oráculo. Ao estabelecer o que é teoricamente possível com a massa de atenção densa, isola as fontes de erro em sistemas esparsos práticos.
2. Definição do Oráculo: O oráculo de massa de atenção top-k serve como uma "referência controlada" para separar a viabilidade do orçamento da performance do indexador e da realização do tempo de execução.
3. Evidência Empírica: Fornece a primeira evidência de que, para os checkpoints da família Qwen, a massa de atenção densa identifica frequentemente um suporte compacto e dependente da query que preserva o comportamento ao nível da tarefa, mesmo em comprimentos extremos (128K–256K).
4. Separação Metodológica: O relatório desvincula explicitamente três alegações distintas:
   • Viabilidade do Oráculo: O suporte esparso é viável (provado pelas linhas do oráculo).
   • Qualidade do Indexador Destilado: Indexadores de backbone congelado podem aproximar este suporte com perda mínima de qualidade (provado pelas linhas destiladas).
   • Folga de Tempo de Execução: O caminho de serviço esparso oferece reduções significativas de TTFT (provado pelas linhas de estresse e linhas de TTFT destilado).

Limitações e Escopo:
Os autores são modestos em suas alegações, observando que isto é uma "primeira versão". Eles declaram explicitamente que:

• Não existe um único "limite de qualidade-latência denso-oráculo-destilado totalmente correspondido" onde todas as métricas são medidas na exata mesma configuração.
• Os indexadores destilados foram validados primariamente em 16K/32K, enquanto as maiores acelerações de TTFT foram medidas em 128K/256K sem verificações de qualidade pareadas.
• O design do indexador de colapso de cabeça é apresentado como uma escolha orientada ao sistema, não como um ótimo provado por ablação contra variantes de expansão de cabeça.
• A evidência está concentrada atualmente na família Qwen; a validação entre famílias de modelos é deixada para trabalhos futuros.

Em resumo, o artigo demonstra que a atenção densa é frequentemente um "excesso" (overkill) para a recuperação de contexto longo em camadas GQA e que um indexador de colapso de cabeça destilado pode aproximar efetivamente o suporte necessário, preservando a qualidade enquanto desbloqueia acelerações significativas de TTFT, desde que o suporte e o orçamento sejam gerenciados cuidadosamente.