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

Abstract

Die Skalierung der Kontextlänge großer Sprachmodelle (LLMs) bietet erhebliche Vorteile, ist jedoch rechenintensiv. Dieser Aufwand resultiert vor allem aus der Selbst-Attention-Mechanismus, dessen O(N²)-Komplexität bezüglich der Sequenzlänge eine wesentliche Engstelle sowohl für Speicher als auch für Latenz darstellt. Glücklicherweise ist die Aufmerksamkeitsmatrix häufig spärlich, insbesondere bei langen Sequenzen, was Optimierungspotenzial eröffnet. Block-sparse Attention hat sich als vielversprechende Lösung erwiesen, indem Sequenzen in Blöcke unterteilt und die Berechnung für einen Teil dieser Blöcke übersprungen wird. Die Wirksamkeit dieses Ansatzes hängt jedoch stark von den zugrundeliegenden Aufmerksamkeitsmustern ab, was zu suboptimaler Block-sparlichkeit führen kann. Beispielsweise können wichtige Key-Token für eine Anfrage innerhalb eines einzelnen Blocks über zahlreiche andere Blöcke verteilt sein, was zu rechnerischer Redundanz führt. In dieser Arbeit stellen wir Permuted Block-Sparse Attention (PBS-Attn) vor, eine plug-and-play-Methode, die die Permutations-Eigenschaften der Aufmerksamkeit nutzt, um die Block-sparlichkeit zu erhöhen und die rechnerische Effizienz der LLM-Prefilling-Phase zu verbessern. Wir führen umfassende Experimente auf anspruchsvollen realen Lang-Kontext-Datensätzen durch und zeigen, dass PBS-Attn die bestehenden Block-sparse-Attention-Methoden in Bezug auf Modellgenauigkeit konsistent übertrifft und sich nahe an die Leistung der vollständigen Attention-Baselinemethode annähert. Durch die Nutzung unserer maßgeschneiderten permuted-FlashAttention-Kerne erreicht PBS-Attn eine end-to-end-Geschwindigkeitssteigerung von bis zu 2,75-fach bei der Verarbeitung langer Kontexte, was seine praktische Relevanz belegt. Der Quellcode ist unter https://github.com/xinghaow99/pbs-attn verfügbar.