CompactAttention: Acelerando o Preenchimento Fragmentado com Seleção de KV por União de Blocos

Resumo

O preenchimento em blocos tornou-se uma estratégia de serviço amplamente adotada para modelos de linguagem de grande porte com contexto longo, mas o cálculo eficiente da atenção nesse regime ainda é desafiador. Os métodos existentes de atenção esparsa são projetados principalmente para preenchimento único e não se traduzem eficientemente para o preenchimento em blocos: kernels esparsos em blocos perdem eficiência quando o comprimento da consulta é limitado pelo tamanho do bloco, enquanto a busca por padrões refinados torna-se custosa quando repetida sobre o cache KV acumulado a cada bloco. O QUOKA, um método recente que visa diretamente o preenchimento em blocos, evita a sobrecarga de kernels esparsos, mas depende de seleção KV em nível de token com subamostragem de consultas, o que pode perder entradas KV específicas da consulta e introduzir sobrecarga explícita de cópia KV. Para superar essas limitações, propomos o CompactAttention, um mecanismo de atenção para preenchimento em blocos baseado em Seleção KV por União de Blocos. O CompactAttention trata máscaras esparsas 2D em blocos como sinais de seleção KV, em vez de planos diretos de execução de kernels esparsos, e as converte em tabelas de blocos KV por grupo cientes de GQA por meio de união de blocos Q e união intragrupo. Essa construção produz as tabelas de blocos mínimas que preservam todos os blocos KV selecionados pelas máscaras de entrada sob restrições de execução paginada, permitindo que os blocos KV selecionados sejam acessados in-place sem compactação KV explícita. No LLaMA-3.1-8B-Instruct, o CompactAttention mantém precisão próxima à atenção densa no benchmark RULER, enquanto oferece até 2,72 vezes de aceleração da atenção em comprimento de contexto de 128K sob preenchimento em blocos.

English

Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging. Existing sparse attention methods are primarily designed for one-shot prefill and do not translate efficiently to chunked prefill: block-sparse kernels lose efficiency when the query length is limited by the chunk size, while fine-grained pattern search becomes costly when repeated over the accumulated KV cache at every chunk. QUOKA, a recent method that directly targets chunked prefill, avoids sparse-kernel overhead but relies on query-subsampled, token-level KV selection, which can miss query-specific KV entries and introduce explicit KV-copy overhead. To address these limitations, we propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection. CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction produces the minimal block tables that preserve all KV blocks selected by the input masks under paged execution constraints, enabling selected KV blocks to be accessed in place without explicit KV compaction. On LLaMA-3.1-8B-Instruct, CompactAttention maintains accuracy close to dense attention on the RULER benchmark while delivering up to 2.72times attention speedup at 128K context length under chunked prefill.