💬 NLP

CoSA: Accelerating Long-Context Inference via Proxy-Kernel Co-Designed Sparse Attention

O CoSA é um framework de atenção esparsa de dois estágios e livre de treinamento que acopla um Proxy Sensível ao Kernel com um Kernel de Pulo Ordenado para otimizar dinamicamente a seleção e o salto de blocos, alcançando aumentos significativos na velocidade de inferência e redução de latência para LLMs de contexto longo, mantendo alta precisão sob orçamentos computacionais restritos.

Autores originais: Yufei Xue, Lin Niu, Hong Liu, Siran Liu, Hanyong Shao, Wei Liu, Guanghua Yu, Jianchen Zhu, Jun Zhang

Publicado 2026-07-29

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Yufei Xue, Lin Niu, Hong Liu, Siran Liu, Hanyong Shao, Wei Liu, Guanghua Yu, Jianchen Zhu, Jun Zhang

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você esteja tentando ler uma enciclopédia massiva de 128.000 páginas para responder a uma única pergunta. No mundo da Inteligência Artificial, essas "enciclopédias" são chamadas de Grandes Modelos de Linguagem (LLMs), e o processo de "leitura" é como elas entendem o contexto. O problema é que a maneira padrão como esses modelos leem é como um bibliotecário que insiste em ler cada página do livro, da primeira palavra até a última, antes mesmo de poder começar a pensar na sua pergunta. À medida que o livro fica mais longo, essa abordagem de "ler tudo" torna-se dolorosamente lenta e cara, como tentar correr uma maratona carregando uma mochila cheia de tijolos.

Para corrigir isso, cientistas tentaram ensinar o bibliotecário a ser um pouco mais seletivo, uma técnica chamada "atenção esparsa" (sparse attention). Em vez de ler todas as páginas, o modelo tenta adivinhar quais páginas são importantes e pula o resto. Geralmente, isso envolve duas etapas: primeiro, um "proxy" rápido (um adivinhador veloz) olha para o livro e marca as páginas importantes com uma lista simples de "Sim" ou "Não". Segundo, um "kernel" (o trabalhador real) segue essa lista e faz o trabalho pesado. É um sistema decente, mas tem uma falha: quando o livro fica enorme e você precisa ser muito rigoroso sobre quais páginas pular para economizar tempo, o adivinhador rápido começa a cometer erros, e o trabalhador apenas segue a lista ruim cegamente. O resultado é que o modelo torna-se mais rápido, mas começa a esquecer detalhes importantes, como um bibliotecário que pula o clímax da história porque achou que era entediante.

É aqui que entra um novo método chamado CoSA (Atenção Esparsa Co-Projetada). Os pesquisadores por trás do CoSA perceberam que o "adivinhador" e o "trabalhador" estavam trabalhando isoladamente, o que fazia o sistema falhar sob pressão. Eles decidiram redesenhar todo o processo para que os dois trabalhem juntos como uma equipe. Em vez de apenas entregar ao trabalhador uma lista simples de "Sim/Não", o novo "adivinhador" (chamado de Kernel-Aware Proxy, ou KAP) entrega uma lista de prioridades. Ele não diz apenas "leia esta página"; ele diz: "Leia esta página primeiro, depois esta, depois aquela".

Aqui está o truque de mágica: o trabalhador (chamado de Ordered-Skipping Kernel, ou OSK) usa essa lista de prioridades para reorganizar a ordem em que lê as páginas. Ao ler as páginas mais críticas primeiro, o trabalhador constrói uma "pontuação contínua" do que é importante muito cedo no processo. Como ele sabe as coisas mais importantes logo de início, pode pular com confiança ainda mais páginas mais adiante no processo sem se confundir. É como se você estivesse lendo um romance de mistério e o detetive lhe dissesse: "Leia os três primeiros capítulos para encontrar o assassino, depois você pode pular com segurança os próximos cinquenta capítulos de dicas de jardinagem".

O artigo mostra que essa abordagem de trabalho em equipe é um divisor de águas. Quando testado em modelos lendo contextos de até 128.000 tokens (aproximadamente o tamanho de um romance longo), o CoSA conseguiu acelerar a parte da "atenção" do processo em 4,93 vezes e reduziu o tempo para gerar a primeira palavra em 2,53 vezes. Crucialmente, ele fez isso sem que o modelo perdesse sua capacidade de entender a história ou resolver problemas complexos de raciocínio. Os pesquisadores descobriram que, ao permitir que o proxy e o kernel conversassem entre si e compartilhassem uma ordem específica de operações, eles podiam ser muito mais agressivos ao pular o trabalho, mantendo o modelo inteligente. Acontece que, no mundo da IA, saber quando ler uma página é tão importante quanto saber qual página ler.

Resumo Técnico: CoSA - Acelerando a Inferência de Longo Contexto via Atenção Esparsa Co-projetada com Kernel de Proxy

1. Definição do Problema

O custo computacional quadrático dos mecanismos de autoatenção torna a inferência de contexto longo proibitivamente cara para os modelos de linguagem de grande escala (LLMs) modernos. Embora a atenção esparsa ofereça um remédio ao computar seletivamente interações importantes de query-key, as abordagens existentes enfrentam um compromisso fundamental entre eficiência e fidelidade, particularmente sob orçamentos de esparsidade rigorosos.

Os métodos atuais geralmente se dividem em duas categorias:

Abordagens Baseadas em Proxy: Utilizam um proxy leve para prever uma máscara esparsa binária antes da execução do kernel. Embora eficazes em orçamentos moderados, esses proxies inevitavelmente perdem blocos salientes à medida que o orçamento diminui, levando a uma degradação significativa da precisão.
Abordagens de Pulo de Kernel (In-Kernel Skipping): Utilizam estatísticas exatas de Softmax dentro do kernel para pular blocos insignificantes. Embora precisas, sofrem de duas limitações:
- Requisito de Proxy Denso: Devem computar a interação completa Query-Key (QK) para cada bloco para tomar uma decisão de salto, falhando em eliminar o custo quadrático dominante de QK.
- Pulo Conservador: Devido ao "efeito balde" (onde uma única linha outlier veta um salto) e ao uso de um máximo móvel (que é um limite inferior do máximo global) em vez de um máximo global, esses métodos aplicam regras de salto excessivamente conservadoras, deixando potenciais ganhos de velocidade inexplorados.

O gargalo central identificado é a desconexão entre o proxy e o kernel de backend. Os proxies geram máscaras binárias estáticas sem considerar a ordem de execução, enquanto os kernels executam em ordens fixas que dificultam a esparsidade ideal.

2. Metodologia: CoSA

Os autores propõem o CoSA, um mecanismo de atenção esparsa por blocos de dois estágios e livre de treinamento, construído sobre o princípio de co-projeto de proxy-kernel. O CoSA acopla um Kernel-Aware Proxy (KAP) a um Ordered-Skipping Kernel (OSK) para otimizar conjuntamente a seleção de blocos e a ordem de execução.

2.1 Princípios de Design Central

Invariância de Ordem: Aproveitando a propriedade de que a acumulação da Online-Softmax (OSM) é independente da ordem em que os blocos de chave são visitados, o CoSA permite o reordenamento arbitrário do processamento de blocos sem alterar o resultado final.
Máscara de Computação-Ordem: Em vez de uma máscara binária, o CoSA gera uma máscara que prescreve a ordem específica em que as páginas de KV devem ser visitadas.

2.2 Estágio 1: Kernel-Aware Proxy (KAP)

O KAP opera sob um orçamento moderado para selecionar blocos e determinar sua ordem de visita.

Ranking de Blocos: Ele realiza o downsampling de queries e keys, computa scores aproximados e aplica MaxPool sobre as chaves seguido de MaxNorm. Este processo preserva o "rowmax" (o valor máximo em uma linha) dentro dos blocos e converte scores em contribuições de ranking relativo.
Geração de Flag HRM: Identifica blocos "Have Rowmax" (HRM) — blocos que contêm o valor máximo para uma linha de query específica.
Ordenação por Grupo de Flags: O KAP constrói uma máscara de ordem de computação ( $M_{KAP}$ ) que prioriza a visita aos blocos HRM, seguidos pelos outros blocos ordenados por seus scores estimados. Esta ordenação garante que o kernel encontre o máximo global precocemente na trajetória.

2.3 Estágio 2: Ordered-Skipping Kernel (OSK)

O OSK consome a máscara $M_{KAP}$ para executar a atenção esparsa com três padrões distintos:

Saltado pela Máscara (Mask-Skipped): Blocos não selecionados pelo KAP são fisicamente saltados via remapeamento de página de KV. Isso elimina os custos de controle de fluxo e sincronização associados aos blocos saltados em kernels esparsos padrão.
Saltado no Kernel (In-Kernel Skipped): Para os blocos retidos pelo Estágio 1, o OSK realiza um segundo estágio de salto usando estatísticas exatas de online-softmax. Como os blocos HRM foram visitados primeiro (devido à ordenação do KAP), o máximo móvel sobe rapidamente para o máximo global. Isso permite que a condição de salto (baseada na diferença entre o máximo local e o máximo móvel) seja satisfeita de forma mais agressiva e segura, evitando o "efeito balde".
Computado: Blocos que passam por ambos os filtros são totalmente computados.

3. Principais Contribuições

Revelação do Desconecte: O artigo identifica que o isolamento entre o design do proxy e a execução do kernel é a barreira fundamental para alcançar alta esparsidade sem perda de precisão.
Design do KAP: Um novo proxy que emite uma máscara de ordem de computação em vez de uma máscara binária. Ele seleciona blocos sob um orçamento moderado e prescreve uma ordem de execução que prioriza blocos contendo valores de rowmax.
Design do OSK: Um kernel otimizado que salta fisicamente páginas baseando-se na máscara do KAP e realiza o salto agressivo no kernel aproveitando as estatísticas de máximo móvel melhoradas derivadas da trajetória reordenada.
Validação Empírica: Experimentos extensos demonstrando que o CoSA alcança maior precisão com orçamentos menores em comparação com métodos de estado da arte (MInference, FlexPrefill, XAttention) em vários backbones de LLM (Qwen3-8B, Llama-3.1-8B-Instruct) e benchmarks.

4. Resultados Experimentais

Os experimentos foram conduzidos no RULER (tarefas sintéticas) e LongBench-v2 (raciocínio no mundo real) com comprimentos de contexto variando de 4K a 128K.

Precisão: O CoSA supera consistentemente os métodos esparsos existentes. No LongBench-v2 com Qwen3-8B, o CoSA alcançou a maior precisão média (36,51%) operando com o menor orçamento (15%). Manteve um desempenho robusto tanto no modo "w/o CoT" (direto) quanto no modo "w/ CoT" (raciocínio).
Eficiência:
- Aceleração de Atenção: O CoSA alcança um ganho de velocidade de 4,93× na computação de atenção em um contexto de 128K.
- Latência End-to-End: Reduz o Tempo para o Primeiro Token (TTFT) em 2,53× em 128K de contexto com degradação de desempenho negligível em comparação com a atenção densa.
Estudos de Ablação: O artigo demonstra que a combinação de KAP (seleção ordenada) e OSK (remapeamento de página + salto no kernel) é sinérgica. Usar apenas o KAP melhora a precisão; adicionar o salto no kernel aumenta a esparsidade; e adicionar o remapeamento de página restaura e melhora a precisão enquanto aperta ainda mais o orçamento.

5. Significância e Alegações

O artigo afirma que o CoSA representa uma mudança de paradigma de tratar o proxy e o kernel como componentes isolados para um sistema co-projetado. Ao uni-los com uma máscara de ordem de computação, o CoSA resolve o compromisso entre eficiência e fidelidade que assola os métodos atuais de atenção esparsa.

Os autores afirmam que:

A esparsidade agressiva é alcançável sem perda de precisão: Ao antecipar os blocos HRM, o kernel pode saltar mais blocos de forma segura, empurrando o orçamento de esparsidade para níveis mais baixos do que anteriormente possível.
A execução física importa: Ter apenas uma boa máscara é insuficiente; o kernel deve reordenar fisicamente o acesso à memória (via remapeamento de página) para realizar os ganhos de velocidade teóricos e evitar custos de sincronização redundantes.
Escalabilidade: O CoSA é particularmente eficaz para contextos longos (128K+), onde o custo quadrático da atenção domina, oferecendo reduções substanciais de latência no mundo real para aplicações como Geração Aumentada de Recuperação (RAG) e agentes autônomos.

O artigo conclui observando que, embora o CoSA foque atualmente na fase de preenchimento (prefilling), trabalhos futuros explorarão a extensão desses princípios de co-projeto para a fase de decodificação, que apresenta padrões computacionais distintos.