CompactAttention : Accélération du préremplissage par blocs avec la sélection KV par union de blocs

Résumé

Le préremplissage par blocs est devenu une stratégie de service largement adoptée pour les grands modèles de langage à contexte long, mais le calcul efficace de l'attention dans ce régime reste difficile. Les méthodes d'attention sparse existantes sont principalement conçues pour le préremplissage en une seule fois et ne se traduisent pas efficacement en préremplissage par blocs : les noyaux blocs-sparses perdent en efficacité lorsque la longueur des requêtes est limitée par la taille du bloc, tandis que la recherche de motifs à grain fin devient coûteuse lorsqu'elle est répétée sur le cache KV accumulé à chaque bloc. QUOKA, une méthode récente ciblant directement le préremplissage par blocs, évite le surcoût des noyaux sparses mais repose sur une sélection KV au niveau des tokens sous-échantillonnée par requête, ce qui peut omettre des entrées KV spécifiques aux requêtes et introduire un surcoût explicite de copie KV. Pour remédier à ces limitations, nous proposons CompactAttention, un mécanisme d'attention pour préremplissage par blocs basé sur la Sélection KV par Union de Blocs. CompactAttention traite les masques sparses 2D par blocs comme des signaux de sélection KV plutôt que comme des plans d'exécution directs de noyaux sparses, et les convertit en tables de blocs KV par groupe tenant compte de GQA via une union de blocs Q et une union intra-groupe. Cette construction produit les tables de blocs minimales qui préservent tous les blocs KV sélectionnés par les masques d'entrée sous contraintes d'exécution paginée, permettant d'accéder aux blocs KV sélectionnés sur place sans compaction KV explicite. Sur LLaMA-3.1-8B-Instruct, CompactAttention maintient une précision proche de celle de l'attention dense sur le benchmark RULER tout en offrant une accélération de l'attention allant jusqu'à 2,72 fois pour une longueur de contexte de 128K en préremplissage par blocs.

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.