Sparsere block-sparse Aufmerksamkeit durch Token-Permutation

papers.abstract

Die Skalierung der Kontextlänge von großen Sprachmodellen (LLMs) bietet erhebliche Vorteile, ist jedoch rechenintensiv. Diese Kosten ergeben sich hauptsächlich aus dem Self-Attention-Mechanismus, dessen O(N²)-Komplexität in Bezug auf die Sequenzlänge einen großen Engpass für Speicher und Latenz darstellt. Glücklicherweise ist die Attention-Matrix oft spärlich besetzt, besonders bei langen Sequenzen, was eine Optimierungsmöglichkeit nahelegt. Block-sparse Attention hat sich als vielversprechende Lösung erwiesen, die Sequenzen in Blöcke unterteilt und die Berechnung für eine Teilmenge dieser Blöcke überspringt. Die Effektivität dieser Methode hängt jedoch stark von den zugrundeliegenden Attention-Mustern ab, was zu suboptimaler sparsity auf Blockebene führen kann. Beispielsweise können wichtige Key-Tokens für Queries innerhalb eines einzelnen Blocks über zahlreiche andere Blöcke verstreut sein, was zu rechentechnischer Redundanz führt. In dieser Arbeit schlagen wir Permuted Block-Sparse Attention (PBS-Attn) vor, eine Plug-and-Play-Methode, die die Permutationseigenschaften von Attention nutzt, um die sparsity auf Blockebene zu erhöhen und die Recheneffizienz der LLM-Prefilling-Phase zu verbessern. Wir führen umfassende Experimente auf anspruchsvollen realen Langkontext-Datensätzen durch, die zeigen, dass PBS-Attn bestehende Block-sparse-Attention-Methoden in der Modellgenauigkeit konsequent übertrifft und sich eng an die Baseline der vollständigen Attention annähert. Gestützt durch unsere maßgeschneiderten permuted-FlashAttention-Kernel erreicht PBS-Attn eine End-to-End-Beschleunigung von bis zu 2,75x beim Langkontext-Prefilling, was seine praktische Tauglichkeit bestätigt. Code verfügbar unter 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 block-sparse Aufmerksamkeit durch Token-Permutation

Sparser Block-Sparse Attention via Token Permutation

papers.abstract

Support