Sparsere Blok-Sparse Aandacht via Tokenpermutatie

Samenvatting

Het vergroten van de contextlengte van grote taalmodel(len) (LLM's) biedt aanzienlijke voordelen, maar is rekenkundig kostbaar. Deze kosten zijn voornamelijk het gevolg van het self-attention-mechanisme, waarvan de O(N²)-complexiteit ten opzichte van de sequentielengte een grote bottleneck vormt voor zowel geheugen als latentie. Gelukkig is de attention-matrix vaak sparse, vooral bij lange sequenties, wat een mogelijkheid tot optimalisatie suggereert. Block-sparse attention is naar voren gekomen als een veelbelovende oplossing die sequenties opdeelt in blokken en de berekening voor een subset van deze blokken overslaat. De effectiviteit van deze methode is echter sterk afhankelijk van de onderliggende attention-patronen, wat kan leiden tot suboptimale sparsity op blokniveau. Belangrijke key-tokens voor queries binnen een enkel blok kunnen bijvoorbeeld verspreid zijn over talrijke andere blokken, wat leidt tot rekenkundige redundantie. In dit werk stellen we Permuted Block-Sparse Attention (PBS-Attn) voor, een plug-and-play-methode die gebruikmaakt van de permutatie-eigenschappen van attention om de sparsity op blokniveau te vergroten en de rekenkundige efficiëntie van LLM-prefilling te verbeteren. We voeren uitgebreide experimenten uit met uitdagende real-world datasets met een lange context, waaruit blijkt dat PBS-Attn consequent beter presteert dan bestaande block-sparse attention-methoden in modelnauwkeurigheid en de baseline van volledige attention dicht benadert. Aangedreven door onze aangepaste permuted-FlashAttention-kernels bereikt PBS-Attn een end-to-end versnelling van tot 2,75x bij prefilling met een lange context, wat de praktische haalbaarheid bevestigt. Code beschikbaar op 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

Sparsere Blok-Sparse Aandacht via Tokenpermutatie

Sparser Block-Sparse Attention via Token Permutation

Samenvatting

Support