CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

Resumo Técnico: CompactAttention

Declaração do Problema
À medida que os Modelos de Linguagem de Grande Escala (LLMs) evoluem para suportar janelas de contexto de centenas de milhares a milhões de tokens, o processamento de sequências longas por meio de uma única passagem de pré-preenchimento "de uma só vez" tornou-se impraticável devido aos custos computacionais quadráticos e ao risco de bloquear solicitações de decodificação em sistemas de serviço online. Consequentemente, o pré-preenchimento fragmentado (chunked prefill)—processar entradas sequencialmente em blocos de tamanho fixo enquanto se atende aos caches de Chave-Valor (KV) acumulados—tornou-se a estratégia de serviço padrão em frameworks como vLLM e SGLang.

No entanto, os métodos de aceleração existentes enfrentam limitações significativas neste regime:

1. Ineficiência da Atenção Esparsa por Blocos: Os métodos tradicionais de esparsidade por blocos (por exemplo, XAttention, SeerAttention) dependem de kernels esparsos que exigem comprimentos de consulta grandes para amortizar a sobrecarga do acesso irregular à memória e da interpretação de máscaras. No pré-preenchimento fragmentado, o comprimento da consulta é limitado pelo tamanho do bloco ($Q \ll KV$), fazendo com que esses kernels percam paralelismo e eficiência, muitas vezes falhando em superar a atenção densa.
2. Sobrecarga de Busca de Padrões: Os métodos esparsos exigem a descoberta de padrões esparsos dependentes da entrada. No pré-preenchimento fragmentado, essa busca deve ser repetida sobre o cache KV acumulado em crescimento para cada bloco, tornando a sobrecarga cumulativa de busca um gargalo crítico.
3. Limitações da Seleção Direta de KV: Métodos como QUOKA evitam kernels esparsos selecionando um subconjunto de tokens KV com base em um conjunto de consultas subamostradas e realizando atenção densa. No entanto, essa abordagem sofre de dois problemas:
   • Lacunas de Cobertura: Entradas KV específicas de consultas importantes para consultas não amostradas podem ser perdidas, levando à degradação da precisão em tarefas que exigem acesso distribuído à informação.
   • Sobrecarga de Cópia: A seleção ao nível de token exige a coleta e cópia explícita de entradas KV em um buffer compacto antes da execução da atenção, introduzindo uma sobrecarga que escala com o comprimento do contexto e o tamanho do lote.

Metodologia: CompactAttention
Os autores propõem o CompactAttention, um mecanismo que desacopla a seleção de blocos KV em nível de bloco da execução da atenção. Em vez de usar máscaras esparsas 2D como planos de execução diretos para kernels esparsos, o CompactAttention trata-as como sinais de seleção de KV para construir tabelas de blocos KV por grupo conscientes de GQA.

O método opera em duas etapas:

1. Etapa de Seleção (Construção da Tabela de Blocos KV de União):

   • O sistema aceita uma máscara esparsa por bloco por cabeça 2D gerada por qualquer método de busca de padrões leve (por exemplo, SeerAttention ou FlashPrefill).
   • Aplica-se uma união de blocos Q através dos blocos de consulta dentro do bloco atual. Isso é necessário porque a atenção paginada densa consome uma única lista de blocos KV para um grupo de blocos de consulta executados juntos, em vez de decisões individuais para cada telha $(Q\text{-bloco}, KV\text{-bloco})$.
   • Aplica-se uma união intra-grupo através das cabeças de consulta que compartilham um grupo KV (padrão para grupos de Atenção de Consulta Agrupada - Grouped-Query Attention).
   • O resultado é uma tabela mínima de blocos KV por grupo ($T_{b,g}$) que preserva todos os blocos KV selecionados pelas máscaras de entrada originais. Essa construção garante a preservação da cobertura: qualquer bloco KV selecionado por qualquer bloco de consulta ou cabeça no grupo permanece selecionado.

2. Etapa de Execução (Atenção Paginada Zero-Cópia):

   • Diferentemente dos métodos de seleção ao nível de token, o CompactAttention opera em granularidade de bloco.
   • Utiliza um layout de memória KV-head-major $[B, H_{kv}, L, D]$ em vez do layout padrão sequencial-major. Isso permite que cada tripla $(lote, KV\text{-cabeça}, \text{bloco})$ seja tratada como uma região de memória contígua.
   • As tabelas de blocos construídas são passadas como metadados (especificamente kv_indptr e kv_indices) para um kernel de atenção densa paginada padrão.
   • O kernel acessa os blocos KV selecionados in-place a partir do cache KV original, sem cópia explícita ou compactação em um novo buffer.

Principais Contribuições

• Desacoplamento de Seleção e Execução: O artigo introduz um design que separa a lógica de quais blocos KV selecionar de como eles são executados, permitindo a reutilização de métodos leves de busca de padrões esparsos enquanto evita as ineficiências dos kernels esparsos.
• Construção de União de Blocos: O uso inovador de uniões de blocos Q e intra-grupo para converter máscaras esparsas 2D em tabelas de blocos executáveis, conscientes de GQA, que são mínimas, mas preservam a cobertura.
• Execução Paginada Zero-Cópia: Ao aproveitar um layout KV-head-major e metadados de atenção paginada, o método elimina a sobrecarga de cópia explícita de KV inerente aos métodos de seleção ao nível de token, como o QUOKA.

Resultados Experimentais
Avaliado em LLaMA-3.1-8B-Instruct e Qwen3-30B-A3B-Instruct-2507 sob pré-preenchimento fragmentado:

• Aceleração: Em GPUs H200 com comprimento de contexto de 128K, o CompactAtenção alcança até 2,72× de aceleração na atenção e 1,96× de aceleração ponta a ponta sobre a linha de base de atenção densa. Isso supera significativamente as linhas de base esparsas por blocos (que frequentemente têm desempenho inferior à atenção densa no regime $Q \ll KV$) e o QUOKA (limitado pela sobrecarga de cópia).
• Precisão: O CompactAttention mantém precisão próxima à da atenção densa nos benchmarks RULER e LongBench V2. Supera o QUOKA, particularmente em tarefas que exigem acesso distribuído à informação (por exemplo, Multi-key NIAH), onde a subamostragem de consultas do QUOKA leva à degradação da precisão.
• Eficiência: Estudos de ablação confirmam que a estratégia de execução paginada zero-cópia é mais eficiente do que tanto a execução de kernel esparso quanto a compactação explícita de KV, mesmo considerando a sobrecarga de construção de metadados.

Significado
O artigo afirma que o principal gargalo para a atenção esparsa no pré-preenchimento fragmentado não é meramente quais blocos KV selecionar, mas como os blocos selecionados são executados. Ao tornar a seleção esparsa compatível com kernels eficientes de atenção densa paginada por meio da construção de união de blocos e execução zero-cópia, o CompactAttention oferece um caminho prático para um serviço de LLM de contexto longo mais rápido. Ele aborda as ineficiências específicas do regime $Q \ll KV$ que dificultaram a adoção de métodos de atenção esparsa existentes em sistemas de pré-preenchimento fragmentado.