Attention Sparse par Blocs via Permutation de Tokens

papers.abstract

L'augmentation de la longueur de contexte des grands modèles de langage (LLM) présente des avantages significatifs mais s'avère coûteuse en calculs. Cette dépense provient principalement du mécanisme d'auto-attention, dont la complexité en O(N²) par rapport à la longueur de la séquence constitue un goulot d'étranglement majeur pour la mémoire et la latence. Heureusement, la matrice d'attention est souvent creuse, particulièrement pour les longues séquences, ce qui suggère une opportunité d'optimisation. L'attention bloc-creuse est apparue comme une solution prometteuse qui partitionne les séquences en blocs et ignore le calcul pour un sous-ensemble de ces blocs. Cependant, l'efficacité de cette méthode dépend fortement des schémas d'attention sous-jacents, qui peuvent entraîner une creusité au niveau des blocs sous-optimale. Par exemple, les jetons clés importants pour des requêtes au sein d'un seul bloc peuvent être dispersés sur de nombreux autres blocs, conduisant à une redondance computationnelle. Dans ce travail, nous proposons Permuted Block-Sparse Attention (PBS-Attn), une méthode prête à l'emploi qui exploite les propriétés de permutation de l'attention pour augmenter la creusité au niveau des blocs et améliorer l'efficacité computationnelle du préremplissage des LLM. Nous menons des expériences complètes sur des ensembles de données réels et exigeants à long contexte, démontrant que PBS-Attn surpasse constamment les méthodes d'attention bloc-creuse existantes en précision du modèle et se rapproche étroitement de la référence d'attention complète. Propulsée par nos noyaux personnalisés permuted-FlashAttention, PBS-Attn atteint une accélération de bout en bout allant jusqu'à 2,75 fois dans le préremplissage à long contexte, confirmant sa viabilité pratique. Code disponible à l'adresse https://github.com/xinghaow99/pbs-attn.

English

Scaling the context length of large language models (LLMs) offers significant benefits but is computationally expensive. This expense stems primarily from the self-attention mechanism, whose O(N^2) complexity with respect to sequence length presents a major bottleneck for both memory and latency. Fortunately, the attention matrix is often sparse, particularly for long sequences, suggesting an opportunity for optimization. Block-sparse attention has emerged as a promising solution that partitions sequences into blocks and skips computation for a subset of these blocks. However, the effectiveness of this method is highly dependent on the underlying attention patterns, which can lead to sub-optimal block-level sparsity. For instance, important key tokens for queries within a single block may be scattered across numerous other blocks, leading to computational redundancy. In this work, we propose Permuted Block-Sparse Attention (PBS-Attn), a plug-and-play method that leverages the permutation properties of attention to increase block-level sparsity and enhance the computational efficiency of LLM prefilling. We conduct comprehensive experiments on challenging real-world long-context datasets, demonstrating that PBS-Attn consistently outperforms existing block-sparse attention methods in model accuracy and closely matches the full attention baseline. Powered by our custom permuted-FlashAttention kernels, PBS-Attn achieves an end-to-end speedup of up to 2.75times in long-context prefilling, confirming its practical viability. Code available at https://github.com/xinghaow99/pbs-attn

Attention Sparse par Blocs via Permutation de Tokens

Sparser Block-Sparse Attention via Token Permutation

papers.abstract

Support