Xinghao Wang Pengyu Wang Dong Zhang Chenkun Tan Shaojun Zhou Zhaoxiang Liu Shiguo Lian Fangxu Liu Kai Song Xipeng Qiu

Résumé

L’augmentation de la longueur de contexte des grands modèles linguistiques (LLM) offre des avantages significatifs, mais s’accompagne d’un coût computationnel élevé. Ce coût provient principalement du mécanisme d’attention auto-attentive, dont la complexité O(N²) en fonction de la longueur de la séquence constitue un goulot d’étranglement majeur en termes de mémoire et de latence. Heureusement, la matrice d’attention est souvent creuse, en particulier pour des séquences longues, ce qui ouvre la voie à des optimisations. L’attention par blocs creuse s’est imposée comme une solution prometteuse, en partitionnant les séquences en blocs et en sautant le calcul pour un sous-ensemble de ces blocs. Toutefois, l’efficacité de cette méthode dépend fortement des motifs d’attention sous-jacents, ce qui peut entraîner une sparsité au niveau des blocs suboptimale. Par exemple, des jetons clés importants pour une requête située dans un même bloc peuvent être répartis sur de nombreux autres blocs, provoquant ainsi une redondance computationnelle. Dans ce travail, nous proposons une méthode plug-and-play appelée Permuted Block-Sparse Attention (PBS-Attn), qui exploite les propriétés de permutation de l’attention afin d’augmenter la sparsité au niveau des blocs et d’améliorer l’efficacité computationnelle du pré-remplissage des LLM. Nous menons des expériences approfondies sur des jeux de données réels complexes à long contexte, démontrant que PBS-Attn surpasse de manière cohérente les méthodes existantes d’attention par blocs creux en précision du modèle, tout en se rapprochant étroitement de la borne de référence de l’attention complète. Grâce à nos noyaux personnalisés permuted-FlashAttention, PBS-Attn atteint un accélération en temps total d’up à 2,75 fois lors du pré-remplissage à long contexte, confirmant ainsi sa viabilité pratique. Le code est disponible à l’adresse : https://github.com/xinghaow99/pbs-attn